1st cut at gradio UI

finetune_gui.py

@@ -0,0 +1,454 @@
+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
+
+
+def save_variables(
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_model,
+    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
+):
+    # Return the values of the variables as a dictionary
+    variables = {
+        "pretrained_model_name_or_path": pretrained_model_name_or_path,
+        "v2": v2,
+        "v_model": v_model,
+        # "model_list": model_list,
+        "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
+    }
+
+    # Save the data to the selected file
+    with open(file_path, "w") as file:
+        json.dump(variables, file)
+
+
+def load_variables(file_path):
+    # load variables from JSON file
+    with open(file_path, "r") as f:
+        my_data = json.load(f)
+
+    # Return the values of the variables as a dictionary
+    return (
+        my_data.get("pretrained_model_name_or_path", None),
+        my_data.get("v2", None),
+        my_data.get("v_model", None),
+        my_data.get("logging_dir", None),
+        # my_data.get("model_list", None),
+        my_data.get("train_data_dir", None),
+        my_data.get("reg_data_dir", None),
+        my_data.get("output_dir", None),
+        my_data.get("max_resolution", None),
+        my_data.get("learning_rate", None),
+        my_data.get("lr_scheduler", None),
+        my_data.get("lr_warmup", None),
+        my_data.get("train_batch_size", None),
+        my_data.get("epoch", None),
+        my_data.get("save_every_n_epochs", None),
+        my_data.get("mixed_precision", None),
+        my_data.get("save_precision", None),
+        my_data.get("seed", None),
+        my_data.get("num_cpu_threads_per_process", None),
+        my_data.get("convert_to_safetensors", None),
+        my_data.get("convert_to_ckpt", None)
+    )
+
+
+def train_model(
+    pretrained_model_name_or_path,
+    v2,
+    v_model,
+    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_input
+):
+    def save_inference_file(output_dir, v2, v_model):
+        # Copy inference model for v2 if required
+        if v2 and v_model:
+            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")
+
+    # calculate max_train_steps
+    max_train_steps = int(
+        math.ceil(float(total_steps) / int(train_batch_size) * int(epoch))
+    )
+    print(f"max_train_steps = {max_train_steps}")
+
+    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_model:
+        run_cmd += " --v_parameterization"
+    if cache_latent_input:
+        run_cmd += " --cache_latents"
+    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}"
+
+    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_model)
+
+        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_model)
+    else:
+        # Copy inference model for v2 if required
+        save_inference_file(output_dir, v2, v_model)
+
+    # Return the values of the variables as a dictionary
+    # return
+
+
+def set_pretrained_model_name_or_path_input(value, v2, v_model):
+    # define a list of substrings to search for
+    substrings_v2 = ["stable-diffusion-2-1-base", "stable-diffusion-2-base"]
+
+    # check if $v2 and $v_model 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_model = False
+        value = "stabilityai/{}".format(value)
+
+        return value, v2, v_model
+
+    # define a list of substrings to search for v-objective
+    substrings_v_model = ["stable-diffusion-2-1", "stable-diffusion-2"]
+
+    # check if $v2 and $v_model are empty and if $pretrained_model_name_or_path contains any of the substrings in the v_model list
+    if str(value) in substrings_v_model:
+        print("SD v2 v_model detected. Setting --v2 parameter and --v_parameterization")
+        v2 = True
+        v_model = True
+        value = "stabilityai/{}".format(value)
+
+        return value, v2, v_model
+
+    if value == "custom":
+        value = ""
+        v2 = False
+        v_model = False
+
+        return value, v2, v_model
+
+
+# Define the output element
+output = gr.outputs.Textbox(label="Values of variables")
+
+interface = gr.Blocks()
+
+with interface:
+    gr.Markdown("Enter kohya finetuner parameter using this interface.")
+    with gr.Accordion("Configuration File Load/Save", open=False):
+        with gr.Row():
+            config_file_name = gr.inputs.Textbox(label="Config file name", default="")
+            b1 = gr.Button("Load config")
+            b2 = gr.Button("Save config")
+    with gr.Tab("model"):
+        # Define the input elements
+        with gr.Row():
+            pretrained_model_name_or_path_input = gr.inputs.Textbox(
+                label="Pretrained model name or path",
+                placeholder="enter the path to custom model or name of pretrained model",
+            )
+            model_list = gr.Dropdown(
+                label="Model Quick Pick",
+                choices=[
+                    "custom",
+                    "stable-diffusion-2-1-base",
+                    "stable-diffusion-2-base",
+                    "stable-diffusion-2-1",
+                    "stable-diffusion-2",
+                ],
+                value="custom",
+            )
+        with gr.Row():
+            v2_input = gr.inputs.Checkbox(label="v2", default=True)
+            v_model_input = gr.inputs.Checkbox(label="v_model", default=False)
+        model_list.change(
+            set_pretrained_model_name_or_path_input,
+            inputs=[model_list, v2_input, v_model_input],
+            outputs=[pretrained_model_name_or_path_input, v2_input, v_model_input],
+        )
+    with gr.Tab("training dataset and output directory"):
+        train_data_dir_input = gr.inputs.Textbox(
+            label="Image folder", placeholder="directory where the training folders containing the images are located"
+        )
+        reg_data_dir_input = gr.inputs.Textbox(
+            label="Regularisation folder", placeholder="directory where where the regularization folders containing the images are located"
+        )
+        output_dir_input = gr.inputs.Textbox(
+            label="Output directory",
+            placeholder="directory to output trained model",
+        )
+        logging_dir_input = gr.inputs.Textbox(
+            label="Logging directory", placeholder="Optional: enable logging and output TensorBoard log to this directory"
+        )
+        max_resolution_input = gr.inputs.Textbox(
+            label="Max resolution", default="512,512"
+        )
+    with gr.Tab("training parameters"):
+        with gr.Row():
+            learning_rate_input = gr.inputs.Textbox(label="Learning rate", default=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.inputs.Textbox(label="LR warmup", default=0)
+        with gr.Row():
+            train_batch_size_input = gr.inputs.Textbox(
+                label="Train batch size", default=1
+            )
+            epoch_input = gr.inputs.Textbox(label="Epoch", default=1)
+        with gr.Row():
+            save_every_n_epochs_input = gr.inputs.Textbox(
+                label="Save every N epochs", default=1
+            )
+            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",
+            )
+        with gr.Row():
+            seed_input = gr.inputs.Textbox(label="Seed", default=1234)
+            num_cpu_threads_per_process_input = gr.inputs.Textbox(
+                label="Number of CPU threads per process", default=4
+            )
+            cache_latent_input = gr.inputs.Checkbox(
+                label="Cache latent", default=True
+            )
+    with gr.Tab("model conveersion"):
+        convert_to_safetensors_input = gr.inputs.Checkbox(
+            label="Convert to SafeTensors", default=False
+        )
+        convert_to_ckpt_input = gr.inputs.Checkbox(
+            label="Convert to CKPT", default=False
+        )
+    
+    b3 = gr.Button("Run")
+
+    b1.click(
+        load_variables,
+        inputs=[config_file_name],
+        outputs=[
+            pretrained_model_name_or_path_input,
+            v2_input,
+            v_model_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
+        ]
+    )
+    
+    b2.click(
+        save_variables,
+        inputs=[
+            config_file_name,
+            pretrained_model_name_or_path_input,
+            v2_input,
+            v_model_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
+        ]
+    )
+    b3.click(
+        train_model,
+        inputs=[
+            pretrained_model_name_or_path_input,
+            v2_input,
+            v_model_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
+        ]
+    )
+
+# Show the interface
+interface.launch()

requirements.txt

@@ -8,4 +8,6 @@ diffusers[torch]==0.9.0
 pytorch_lightning
 bitsandbytes==0.35.0
 tensorboard
-safetensors==0.2.5
+safetensors==0.2.5
+gradio
+altair