KohyaSS/dreambooth_gui.py

# v1: initial release
# v2: add open and save folder icons

import gradio as gr
import json
import math
import os
import subprocess
import pathlib
import shutil
from glob import glob
from os.path import join
from easygui import fileopenbox, filesavebox, enterbox, diropenbox, msgbox


def save_variables(
    file_path,
    pretrained_model_name_or_path,
    v2,
    v_parameterization,
    logging_dir,
    train_data_dir,
    reg_data_dir,
    output_dir,
    max_resolution,
    learning_rate,
    lr_scheduler,
    lr_warmup,
    train_batch_size,
    epoch,
    save_every_n_epochs,
    mixed_precision,
    save_precision,
    seed,
    num_cpu_threads_per_process,
    convert_to_safetensors,
    convert_to_ckpt,
    cache_latent,
    caption_extention,
    use_safetensors,
    enable_bucket,
    gradient_checkpointing,
    full_fp16,
    no_token_padding,
    stop_text_encoder_training,
    use_8bit_adam,
    xformers,
):
    original_file_path = file_path

    if file_path == None or file_path == "":
        file_path = filesavebox(
            "Select the config file to save",
            default="finetune.json",
            filetypes="*.json",
        )

    if file_path == None:
        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
        return file_path

    # Return the values of the variables as a dictionary
    variables = {
        "pretrained_model_name_or_path": pretrained_model_name_or_path,
        "v2": v2,
        "v_parameterization": v_parameterization,
        "logging_dir": logging_dir,
        "train_data_dir": train_data_dir,
        "reg_data_dir": reg_data_dir,
        "output_dir": output_dir,
        "max_resolution": max_resolution,
        "learning_rate": learning_rate,
        "lr_scheduler": lr_scheduler,
        "lr_warmup": lr_warmup,
        "train_batch_size": train_batch_size,
        "epoch": epoch,
        "save_every_n_epochs": save_every_n_epochs,
        "mixed_precision": mixed_precision,
        "save_precision": save_precision,
        "seed": seed,
        "num_cpu_threads_per_process": num_cpu_threads_per_process,
        "convert_to_safetensors": convert_to_safetensors,
        "convert_to_ckpt": convert_to_ckpt,
        "cache_latent": cache_latent,
        "caption_extention": caption_extention,
        "use_safetensors": use_safetensors,
        "enable_bucket": enable_bucket,
        "gradient_checkpointing": gradient_checkpointing,
        "full_fp16": full_fp16,
        "no_token_padding": no_token_padding,
        "stop_text_encoder_training": stop_text_encoder_training,
        "use_8bit_adam": use_8bit_adam,
        "xformers": xformers,
    }

    # Save the data to the selected file
    with open(file_path, "w") as file:
        json.dump(variables, file)
    
    return file_path


def save_as_variables(
    file_path,
    pretrained_model_name_or_path,
    v2,
    v_parameterization,
    logging_dir,
    train_data_dir,
    reg_data_dir,
    output_dir,
    max_resolution,
    learning_rate,
    lr_scheduler,
    lr_warmup,
    train_batch_size,
    epoch,
    save_every_n_epochs,
    mixed_precision,
    save_precision,
    seed,
    num_cpu_threads_per_process,
    convert_to_safetensors,
    convert_to_ckpt,
    cache_latent,
    caption_extention,
    use_safetensors,
    enable_bucket,
    gradient_checkpointing,
    full_fp16,
    no_token_padding,
    stop_text_encoder_training,
    use_8bit_adam,
    xformers,
):
    original_file_path = file_path

    file_path = filesavebox(
        "Select the config file to save", default="finetune.json", filetypes="*.json"
    )

    if file_path == None:
        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
        return file_path

    # Return the values of the variables as a dictionary
    variables = {
        "pretrained_model_name_or_path": pretrained_model_name_or_path,
        "v2": v2,
        "v_parameterization": v_parameterization,
        "logging_dir": logging_dir,
        "train_data_dir": train_data_dir,
        "reg_data_dir": reg_data_dir,
        "output_dir": output_dir,
        "max_resolution": max_resolution,
        "learning_rate": learning_rate,
        "lr_scheduler": lr_scheduler,
        "lr_warmup": lr_warmup,
        "train_batch_size": train_batch_size,
        "epoch": epoch,
        "save_every_n_epochs": save_every_n_epochs,
        "mixed_precision": mixed_precision,
        "save_precision": save_precision,
        "seed": seed,
        "num_cpu_threads_per_process": num_cpu_threads_per_process,
        "convert_to_safetensors": convert_to_safetensors,
        "convert_to_ckpt": convert_to_ckpt,
        "cache_latent": cache_latent,
        "caption_extention": caption_extention,
        "use_safetensors": use_safetensors,
        "enable_bucket": enable_bucket,
        "gradient_checkpointing": gradient_checkpointing,
        "full_fp16": full_fp16,
        "no_token_padding": no_token_padding,
        "stop_text_encoder_training": stop_text_encoder_training,
        "use_8bit_adam": use_8bit_adam,
        "xformers": xformers,
    }

    # Save the data to the selected file
    with open(file_path, "w") as file:
        json.dump(variables, file)

    return file_path


def open_config_file(
    file_path,
    pretrained_model_name_or_path,
    v2,
    v_parameterization,
    logging_dir,
    train_data_dir,
    reg_data_dir,
    output_dir,
    max_resolution,
    learning_rate,
    lr_scheduler,
    lr_warmup,
    train_batch_size,
    epoch,
    save_every_n_epochs,
    mixed_precision,
    save_precision,
    seed,
    num_cpu_threads_per_process,
    convert_to_safetensors,
    convert_to_ckpt,
    cache_latent,
    caption_extention,
    use_safetensors,
    enable_bucket,
    gradient_checkpointing,
    full_fp16,
    no_token_padding,
    stop_text_encoder_training,
    use_8bit_adam,
    xformers,
):

    original_file_path = file_path
    file_path = get_file_path(file_path)

    if file_path != "" and file_path != None:
        print(file_path)
        # load variables from JSON file
        with open(file_path, "r") as f:
            my_data = json.load(f)
    else:
        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
        my_data = {}

    # Return the values of the variables as a dictionary
    return (
        file_path,
        my_data.get("pretrained_model_name_or_path", pretrained_model_name_or_path),
        my_data.get("v2", v2),
        my_data.get("v_parameterization", v_parameterization),
        my_data.get("logging_dir", logging_dir),
        my_data.get("train_data_dir", train_data_dir),
        my_data.get("reg_data_dir", reg_data_dir),
        my_data.get("output_dir", output_dir),
        my_data.get("max_resolution", max_resolution),
        my_data.get("learning_rate", learning_rate),
        my_data.get("lr_scheduler", lr_scheduler),
        my_data.get("lr_warmup", lr_warmup),
        my_data.get("train_batch_size", train_batch_size),
        my_data.get("epoch", epoch),
        my_data.get("save_every_n_epochs", save_every_n_epochs),
        my_data.get("mixed_precision", mixed_precision),
        my_data.get("save_precision", save_precision),
        my_data.get("seed", seed),
        my_data.get("num_cpu_threads_per_process", num_cpu_threads_per_process),
        my_data.get("convert_to_safetensors", convert_to_safetensors),
        my_data.get("convert_to_ckpt", convert_to_ckpt),
        my_data.get("cache_latent", cache_latent),
        my_data.get("caption_extention", caption_extention),
        my_data.get("use_safetensors", use_safetensors),
        my_data.get("enable_bucket", enable_bucket),
        my_data.get("gradient_checkpointing", gradient_checkpointing),
        my_data.get("full_fp16", full_fp16),
        my_data.get("no_token_padding", no_token_padding),
        my_data.get("stop_text_encoder_training", stop_text_encoder_training),
        my_data.get("use_8bit_adam", use_8bit_adam),
        my_data.get("xformers", xformers),
    )


def train_model(
    pretrained_model_name_or_path,
    v2,
    v_parameterization,
    logging_dir,
    train_data_dir,
    reg_data_dir,
    output_dir,
    max_resolution,
    learning_rate,
    lr_scheduler,
    lr_warmup,
    train_batch_size,
    epoch,
    save_every_n_epochs,
    mixed_precision,
    save_precision,
    seed,
    num_cpu_threads_per_process,
    convert_to_safetensors,
    convert_to_ckpt,
    cache_latent,
    caption_extention,
    use_safetensors,
    enable_bucket,
    gradient_checkpointing,
    full_fp16,
    no_token_padding,
    stop_text_encoder_training_pct,
    use_8bit_adam,
    xformers,
):
    def save_inference_file(output_dir, v2, v_parameterization):
        # Copy inference model for v2 if required
        if v2 and v_parameterization:
            print(f"Saving v2-inference-v.yaml as {output_dir}/last.yaml")
            shutil.copy(
                f"./v2_inference/v2-inference-v.yaml",
                f"{output_dir}/last.yaml",
            )
        elif v2:
            print(f"Saving v2-inference.yaml as {output_dir}/last.yaml")
            shutil.copy(
                f"./v2_inference/v2-inference.yaml",
                f"{output_dir}/last.yaml",
            )

    # Get a list of all subfolders in train_data_dir
    subfolders = [
        f
        for f in os.listdir(train_data_dir)
        if os.path.isdir(os.path.join(train_data_dir, f))
    ]

    total_steps = 0

    # Loop through each subfolder and extract the number of repeats
    for folder in subfolders:
        # Extract the number of repeats from the folder name
        repeats = int(folder.split("_")[0])

        # Count the number of images in the folder
        num_images = len(
            [
                f
                for f in os.listdir(os.path.join(train_data_dir, folder))
                if f.endswith(".jpg")
                or f.endswith(".jpeg")
                or f.endswith(".png")
                or f.endswith(".webp")
            ]
        )

        # Calculate the total number of steps for this folder
        steps = repeats * num_images
        total_steps += steps

        # Print the result
        print(f"Folder {folder}: {steps} steps")

    # Print the result
    # print(f"{total_steps} total steps")

    if reg_data_dir == "":
        reg_factor = 1
    else:
        print(
            "Regularisation images are used... Will double the number of steps required..."
        )
        reg_factor = 2

    # calculate max_train_steps
    max_train_steps = int(
        math.ceil(
            float(total_steps) / int(train_batch_size) * int(epoch) * int(reg_factor)
        )
    )
    print(f"max_train_steps = {max_train_steps}")

    # calculate stop encoder training
    if stop_text_encoder_training_pct == None:
        stop_text_encoder_training = 0
    else:
        stop_text_encoder_training = math.ceil(
            float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
        )
    print(f"stop_text_encoder_training = {stop_text_encoder_training}")

    lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
    print(f"lr_warmup_steps = {lr_warmup_steps}")

    run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_db_fixed.py"'
    if v2:
        run_cmd += " --v2"
    if v_parameterization:
        run_cmd += " --v_parameterization"
    if cache_latent:
        run_cmd += " --cache_latents"
    if use_safetensors:
        run_cmd += " --use_safetensors"
    if enable_bucket:
        run_cmd += " --enable_bucket"
    if gradient_checkpointing:
        run_cmd += " --gradient_checkpointing"
    if full_fp16:
        run_cmd += " --full_fp16"
    if no_token_padding:
        run_cmd += " --no_token_padding"
    if use_8bit_adam:
        run_cmd += " --use_8bit_adam"
    if xformers:
        run_cmd += " --xformers"
    run_cmd += f" --pretrained_model_name_or_path={pretrained_model_name_or_path}"
    run_cmd += f" --train_data_dir={train_data_dir}"
    run_cmd += f" --reg_data_dir={reg_data_dir}"
    run_cmd += f" --resolution={max_resolution}"
    run_cmd += f" --output_dir={output_dir}"
    run_cmd += f" --train_batch_size={train_batch_size}"
    run_cmd += f" --learning_rate={learning_rate}"
    run_cmd += f" --lr_scheduler={lr_scheduler}"
    run_cmd += f" --lr_warmup_steps={lr_warmup_steps}"
    run_cmd += f" --max_train_steps={max_train_steps}"
    run_cmd += f" --use_8bit_adam"
    run_cmd += f" --xformers"
    run_cmd += f" --mixed_precision={mixed_precision}"
    run_cmd += f" --save_every_n_epochs={save_every_n_epochs}"
    run_cmd += f" --seed={seed}"
    run_cmd += f" --save_precision={save_precision}"
    run_cmd += f" --logging_dir={logging_dir}"
    run_cmd += f" --caption_extention={caption_extention}"
    run_cmd += f" --stop_text_encoder_training={stop_text_encoder_training}"

    print(run_cmd)
    # Run the command
    subprocess.run(run_cmd)

    # check if output_dir/last is a directory... therefore it is a diffuser model
    last_dir = pathlib.Path(f"{output_dir}/last")
    print(last_dir)
    if last_dir.is_dir():
        if convert_to_ckpt:
            print(f"Converting diffuser model {last_dir} to {last_dir}.ckpt")
            os.system(
                f"python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.ckpt --{save_precision}"
            )

            save_inference_file(output_dir, v2, v_parameterization)

        if convert_to_safetensors:
            print(f"Converting diffuser model {last_dir} to {last_dir}.safetensors")
            os.system(
                f"python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.safetensors --{save_precision}"
            )

            save_inference_file(output_dir, v2, v_parameterization)
    else:
        # Copy inference model for v2 if required
        save_inference_file(output_dir, v2, v_parameterization)

    # Return the values of the variables as a dictionary
    # return


def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
    # define a list of substrings to search for
    substrings_v2 = [
        "stabilityai/stable-diffusion-2-1-base",
        "stabilityai/stable-diffusion-2-base",
    ]

    # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v2 list
    if str(value) in substrings_v2:
        print("SD v2 model detected. Setting --v2 parameter")
        v2 = True
        v_parameterization = False

        return value, v2, v_parameterization

    # define a list of substrings to search for v-objective
    substrings_v_parameterization = [
        "stabilityai/stable-diffusion-2-1",
        "stabilityai/stable-diffusion-2",
    ]

    # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v_parameterization list
    if str(value) in substrings_v_parameterization:
        print(
            "SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization"
        )
        v2 = True
        v_parameterization = True

        return value, v2, v_parameterization

    # define a list of substrings to v1.x
    substrings_v1_model = [
        "CompVis/stable-diffusion-v1-4",
        "runwayml/stable-diffusion-v1-5",
    ]

    if str(value) in substrings_v1_model:
        v2 = False
        v_parameterization = False

        return value, v2, v_parameterization

    if value == "custom":
        value = ""
        v2 = False
        v_parameterization = False

        return value, v2, v_parameterization


def remove_doublequote(file_path):
    if file_path != None:
        file_path = file_path.replace('"', "")

    return file_path


def get_file_path(file_path):
    file_path = fileopenbox(
        "Select the config file to load", default=file_path, filetypes="*.json"
    )

    return file_path


def get_folder_path():
    folder_path = diropenbox("Select the directory to use")

    return folder_path


css = ""

if os.path.exists("./style.css"):
    with open(os.path.join("./style.css"), "r", encoding="utf8") as file:
        print("Load CSS...")
        css += file.read() + "\n"

interface = gr.Blocks(css=css)

with interface:
    gr.Markdown("Enter kohya finetuner parameter using this interface.")
    with gr.Accordion("Configuration File Load/Save", open=False):
        with gr.Row():
            button_open_config = gr.Button("Open 📂", elem_id="open_folder")
            button_save_config = gr.Button("Save 💾", elem_id="open_folder")
            button_save_as_config = gr.Button("Save as... 💾", elem_id="open_folder")
        config_file_name = gr.Textbox(
            label="", placeholder="type config file path or use buttons..."
        )
        config_file_name.change(
            remove_doublequote, inputs=[config_file_name], outputs=[config_file_name]
        )
    with gr.Tab("Source model"):
        # Define the input elements
        with gr.Row():
            pretrained_model_name_or_path_input = gr.Textbox(
                label="Pretrained model name or path",
                placeholder="enter the path to custom model or name of pretrained model",
            )
            model_list = gr.Dropdown(
                label="(Optional) Model Quick Pick",
                choices=[
                    "custom",
                    "stabilityai/stable-diffusion-2-1-base",
                    "stabilityai/stable-diffusion-2-base",
                    "stabilityai/stable-diffusion-2-1",
                    "stabilityai/stable-diffusion-2",
                    "runwayml/stable-diffusion-v1-5",
                    "CompVis/stable-diffusion-v1-4",
                ],
            )
        with gr.Row():
            v2_input = gr.Checkbox(label="v2", value=True)
            v_parameterization_input = gr.Checkbox(
                label="v_parameterization", value=False
            )
        pretrained_model_name_or_path_input.change(
            remove_doublequote,
            inputs=[pretrained_model_name_or_path_input],
            outputs=[pretrained_model_name_or_path_input],
        )
        model_list.change(
            set_pretrained_model_name_or_path_input,
            inputs=[model_list, v2_input, v_parameterization_input],
            outputs=[
                pretrained_model_name_or_path_input,
                v2_input,
                v_parameterization_input,
            ],
        )

    with gr.Tab("Directories"):
        with gr.Row():
            train_data_dir_input = gr.Textbox(
                label="Image folder",
                placeholder="Directory where the training folders containing the images are located",
            )
            train_data_dir_input_folder = gr.Button("📂", elem_id="open_folder_small")
            train_data_dir_input_folder.click(get_folder_path, outputs=train_data_dir_input)
            reg_data_dir_input = gr.Textbox(
                label="Regularisation folder",
                placeholder="(Optional) Directory where where the regularization folders containing the images are located",
            )
            reg_data_dir_input_folder = gr.Button("📂", elem_id="open_folder_small")
            reg_data_dir_input_folder.click(get_folder_path, outputs=reg_data_dir_input)
        with gr.Row():
            output_dir_input = gr.Textbox(
                label="Output directory",
                placeholder="Directory to output trained model",
            )
            output_dir_input_folder = gr.Button("📂", elem_id="open_folder_small")
            output_dir_input_folder.click(get_folder_path, outputs=output_dir_input)
            logging_dir_input = gr.Textbox(
                label="Logging directory",
                placeholder="Optional: enable logging and output TensorBoard log to this directory",
            )
            logging_dir_input_folder = gr.Button("📂", elem_id="open_folder_small")
            logging_dir_input_folder.click(get_folder_path, outputs=logging_dir_input)
        train_data_dir_input.change(
            remove_doublequote,
            inputs=[train_data_dir_input],
            outputs=[train_data_dir_input],
        )
        reg_data_dir_input.change(
            remove_doublequote,
            inputs=[reg_data_dir_input],
            outputs=[reg_data_dir_input],
        )
        output_dir_input.change(
            remove_doublequote, inputs=[output_dir_input], outputs=[output_dir_input]
        )
        logging_dir_input.change(
            remove_doublequote, inputs=[logging_dir_input], outputs=[logging_dir_input]
        )
    with gr.Tab("Training parameters"):
        with gr.Row():
            learning_rate_input = gr.Textbox(label="Learning rate", value=1e-6)
            lr_scheduler_input = gr.Dropdown(
                label="LR Scheduler",
                choices=[
                    "constant",
                    "constant_with_warmup",
                    "cosine",
                    "cosine_with_restarts",
                    "linear",
                    "polynomial",
                ],
                value="constant",
            )
            lr_warmup_input = gr.Textbox(label="LR warmup", value=0)
        with gr.Row():
            train_batch_size_input = gr.Slider(
                minimum=1, maximum=32, label="Train batch size", value=1, step=1
            )
            epoch_input = gr.Textbox(label="Epoch", value=1)
            save_every_n_epochs_input = gr.Textbox(label="Save every N epochs", value=1)
        with gr.Row():
            mixed_precision_input = gr.Dropdown(
                label="Mixed precision",
                choices=[
                    "no",
                    "fp16",
                    "bf16",
                ],
                value="fp16",
            )
            save_precision_input = gr.Dropdown(
                label="Save precision",
                choices=[
                    "float",
                    "fp16",
                    "bf16",
                ],
                value="fp16",
            )
            num_cpu_threads_per_process_input = gr.Slider(
                minimum=1,
                maximum=os.cpu_count(),
                step=1,
                label="Number of CPU threads per process",
                value=os.cpu_count(),
            )
        with gr.Row():
            seed_input = gr.Textbox(label="Seed", value=1234)
            max_resolution_input = gr.Textbox(label="Max resolution", value="512,512")
        with gr.Row():
            caption_extention_input = gr.Textbox(
                label="Caption Extension",
                placeholder="(Optional) Extension for caption files. default: .caption",
            )
            stop_text_encoder_training_input = gr.Slider(
                minimum=0,
                maximum=100,
                value=0,
                step=1,
                label="Stop text encoder training",
            )
        with gr.Row():
            use_safetensors_input = gr.Checkbox(
                label="Use safetensor when saving", value=False
            )
            enable_bucket_input = gr.Checkbox(label="Enable buckets", value=False)
            cache_latent_input = gr.Checkbox(label="Cache latent", value=True)
            gradient_checkpointing_input = gr.Checkbox(
                label="Gradient checkpointing", value=False
            )
        with gr.Row():
            full_fp16_input = gr.Checkbox(
                label="Full fp16 training (experimental)", value=False
            )
            no_token_padding_input = gr.Checkbox(label="No tokan padding", value=False)
            use_8bit_adam_input = gr.Checkbox(label="Use 8bit adam", value=True)
            xformers_input = gr.Checkbox(label="USe xformers", value=True)

    with gr.Tab("Model conversion"):
        convert_to_safetensors_input = gr.Checkbox(
            label="Convert to SafeTensors", value=False
        )
        convert_to_ckpt_input = gr.Checkbox(label="Convert to CKPT", value=False)

    button_run = gr.Button("Run")

    button_open_config.click(
        open_config_file,
        inputs=[
            config_file_name,
            pretrained_model_name_or_path_input,
            v2_input,
            v_parameterization_input,
            logging_dir_input,
            train_data_dir_input,
            reg_data_dir_input,
            output_dir_input,
            max_resolution_input,
            learning_rate_input,
            lr_scheduler_input,
            lr_warmup_input,
            train_batch_size_input,
            epoch_input,
            save_every_n_epochs_input,
            mixed_precision_input,
            save_precision_input,
            seed_input,
            num_cpu_threads_per_process_input,
            convert_to_safetensors_input,
            convert_to_ckpt_input,
            cache_latent_input,
            caption_extention_input,
            use_safetensors_input,
            enable_bucket_input,
            gradient_checkpointing_input,
            full_fp16_input,
            no_token_padding_input,
            stop_text_encoder_training_input,
            use_8bit_adam_input,
            xformers_input,
        ],
        outputs=[
            config_file_name,
            pretrained_model_name_or_path_input,
            v2_input,
            v_parameterization_input,
            logging_dir_input,
            train_data_dir_input,
            reg_data_dir_input,
            output_dir_input,
            max_resolution_input,
            learning_rate_input,
            lr_scheduler_input,
            lr_warmup_input,
            train_batch_size_input,
            epoch_input,
            save_every_n_epochs_input,
            mixed_precision_input,
            save_precision_input,
            seed_input,
            num_cpu_threads_per_process_input,
            convert_to_safetensors_input,
            convert_to_ckpt_input,
            cache_latent_input,
            caption_extention_input,
            use_safetensors_input,
            enable_bucket_input,
            gradient_checkpointing_input,
            full_fp16_input,
            no_token_padding_input,
            stop_text_encoder_training_input,
            use_8bit_adam_input,
            xformers_input,
        ],
    )

    button_save_config.click(
        save_variables,
        inputs=[
            config_file_name,
            pretrained_model_name_or_path_input,
            v2_input,
            v_parameterization_input,
            logging_dir_input,
            train_data_dir_input,
            reg_data_dir_input,
            output_dir_input,
            max_resolution_input,
            learning_rate_input,
            lr_scheduler_input,
            lr_warmup_input,
            train_batch_size_input,
            epoch_input,
            save_every_n_epochs_input,
            mixed_precision_input,
            save_precision_input,
            seed_input,
            num_cpu_threads_per_process_input,
            convert_to_safetensors_input,
            convert_to_ckpt_input,
            cache_latent_input,
            caption_extention_input,
            use_safetensors_input,
            enable_bucket_input,
            gradient_checkpointing_input,
            full_fp16_input,
            no_token_padding_input,
            stop_text_encoder_training_input,
            use_8bit_adam_input,
            xformers_input,
        ],
        outputs=[config_file_name],
    )

    button_save_as_config.click(
        save_as_variables,
        inputs=[
            config_file_name,
            pretrained_model_name_or_path_input,
            v2_input,
            v_parameterization_input,
            logging_dir_input,
            train_data_dir_input,
            reg_data_dir_input,
            output_dir_input,
            max_resolution_input,
            learning_rate_input,
            lr_scheduler_input,
            lr_warmup_input,
            train_batch_size_input,
            epoch_input,
            save_every_n_epochs_input,
            mixed_precision_input,
            save_precision_input,
            seed_input,
            num_cpu_threads_per_process_input,
            convert_to_safetensors_input,
            convert_to_ckpt_input,
            cache_latent_input,
            caption_extention_input,
            use_safetensors_input,
            enable_bucket_input,
            gradient_checkpointing_input,
            full_fp16_input,
            no_token_padding_input,
            stop_text_encoder_training_input,
            use_8bit_adam_input,
            xformers_input,
        ],
        outputs=[config_file_name],
    )

    button_run.click(
        train_model,
        inputs=[
            pretrained_model_name_or_path_input,
            v2_input,
            v_parameterization_input,
            logging_dir_input,
            train_data_dir_input,
            reg_data_dir_input,
            output_dir_input,
            max_resolution_input,
            learning_rate_input,
            lr_scheduler_input,
            lr_warmup_input,
            train_batch_size_input,
            epoch_input,
            save_every_n_epochs_input,
            mixed_precision_input,
            save_precision_input,
            seed_input,
            num_cpu_threads_per_process_input,
            convert_to_safetensors_input,
            convert_to_ckpt_input,
            cache_latent_input,
            caption_extention_input,
            use_safetensors_input,
            enable_bucket_input,
            gradient_checkpointing_input,
            full_fp16_input,
            no_token_padding_input,
            stop_text_encoder_training_input,
            use_8bit_adam_input,
            xformers_input,
        ],
    )

# Show the interface
interface.launch()