Adding support for new parameters

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.linting.enabled": true
+}

dreambooth_gui.py

@@ -38,7 +38,11 @@ def save_variables(
     use_safetensors,
     enable_bucket,
     gradient_checkpointing,
-    full_fp16
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    use_8bit_adam,
+    xformers,
 ):
     # Return the values of the variables as a dictionary
     variables = {
@@ -67,7 +71,11 @@ def save_variables(
         "use_safetensors": use_safetensors,
         "enable_bucket": enable_bucket,
         "gradient_checkpointing": gradient_checkpointing,
-        "full_fp16": full_fp16
+        "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
@@ -108,6 +116,10 @@ def load_variables(file_path):
         my_data.get("enable_bucket", None),
         my_data.get("gradient_checkpointing", None),
         my_data.get("full_fp16", None),
+        my_data.get("no_token_padding", None),
+        my_data.get("stop_text_encoder_training", None),
+        my_data.get("use_8bit_adam", None),
+        my_data.get("xformers", None),
     )
 
 
@@ -137,7 +149,11 @@ def train_model(
     use_safetensors,
     enable_bucket,
     gradient_checkpointing,
-    full_fp16
+    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
@@ -155,8 +171,11 @@ def train_model(
             )
 
     # 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))]
+    subfolders = [
+        f
+        for f in os.listdir(train_data_dir)
+        if os.path.isdir(os.path.join(train_data_dir, f))
+    ]
 
     total_steps = 0
 
@@ -166,8 +185,16 @@ def train_model(
         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")])
+        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
@@ -182,15 +209,28 @@ def train_model(
     if reg_data_dir == "":
         reg_factor = 1
     else:
-        print("Regularisation images are used... Will double the number of steps required...")
+        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))
+        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}")
 
@@ -209,6 +249,12 @@ def train_model(
         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}"
@@ -227,6 +273,7 @@ def train_model(
     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
@@ -245,8 +292,7 @@ def train_model(
             save_inference_file(output_dir, v2, v_parameterization)
 
         if convert_to_safetensors:
-            print(
-                f"Converting diffuser model {last_dir} to {last_dir}.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}"
             )
@@ -262,7 +308,10 @@ def train_model(
 
 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"]
+    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:
@@ -273,18 +322,26 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
         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"]
+    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")
+        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"]
+    substrings_v1_model = [
+        "CompVis/stable-diffusion-v1-4",
+        "runwayml/stable-diffusion-v1-5",
+    ]
 
     if str(value) in substrings_v1_model:
         v2 = False
@@ -299,16 +356,29 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
 
         return value, v2, v_parameterization
 
+def remove_doublequote(file_path):
+    if file_path != None:
+        file_path = file_path.replace('"', '')
+
+    return file_path
+
+
 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.Textbox(
-                label="Config file name")
+            config_file_name = gr.Textbox(label="Config file name")
             button_load_config = gr.Button("Load config")
             button_save_config = gr.Button("Save config")
+        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():
@@ -325,39 +395,81 @@ with interface:
                     "stabilityai/stable-diffusion-2-1",
                     "stabilityai/stable-diffusion-2",
                     "runwayml/stable-diffusion-v1-5",
-                    "CompVis/stable-diffusion-v1-4"
+                    "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)
+            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],
+            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"
+                label="Image folder",
+                placeholder="Directory where the training folders containing the images are located",
             )
             reg_data_dir_input = gr.Textbox(
-                label="Regularisation folder", placeholder="(Optional) Directory where where the regularization folders containing the images are located"
+                label="Regularisation folder",
+                placeholder="(Optional) Directory where where the regularization folders containing the images are located",
             )
-
         with gr.Row():
             output_dir_input = gr.Textbox(
                 label="Output directory",
                 placeholder="Directory to output trained model",
             )
             logging_dir_input = gr.Textbox(
-                label="Logging directory", placeholder="Optional: enable logging and output TensorBoard log to this directory"
+                label="Logging directory",
+                placeholder="Optional: enable logging and output TensorBoard log to this directory",
             )
+        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)
+            learning_rate_input = gr.Textbox(label="Learning rate", value=1e-6)
             lr_scheduler_input = gr.Dropdown(
                 label="LR Scheduler",
                 choices=[
@@ -372,13 +484,11 @@ with interface:
             )
             lr_warmup_input = gr.Textbox(label="LR warmup", value=0)
         with gr.Row():
-            train_batch_size_input = gr.Textbox(
-                label="Train batch size", value=1
+            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
-            )
+            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",
@@ -398,41 +508,46 @@ with interface:
                 ],
                 value="fp16",
             )
-            num_cpu_threads_per_process_input = gr.Textbox(
-                label="Number of CPU threads per process", value=4
+            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"
-            )
+            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")
-
+                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
-            )
+            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
-        )
+        convert_to_ckpt_input = gr.Checkbox(label="Convert to CKPT", value=False)
 
     button_run = gr.Button("Run")
 
@@ -465,8 +580,12 @@ with interface:
             use_safetensors_input,
             enable_bucket_input,
             gradient_checkpointing_input,
-            full_fp16_input
-        ]
+            full_fp16_input,
+            no_token_padding_input,
+            stop_text_encoder_training_input,
+            use_8bit_adam_input,
+            xformers_input,
+        ],
     )
 
     button_save_config.click(
@@ -498,8 +617,12 @@ with interface:
             use_safetensors_input,
             enable_bucket_input,
             gradient_checkpointing_input,
-            full_fp16_input
-        ]
+            full_fp16_input,
+            no_token_padding_input,
+            stop_text_encoder_training_input,
+            use_8bit_adam_input,
+            xformers_input,
+        ],
     )
     button_run.click(
         train_model,
@@ -529,8 +652,12 @@ with interface:
             use_safetensors_input,
             enable_bucket_input,
             gradient_checkpointing_input,
-            full_fp16_input
-        ]
+            full_fp16_input,
+            no_token_padding_input,
+            stop_text_encoder_training_input,
+            use_8bit_adam_input,
+            xformers_input,
+        ],
     )
 
 # Show the interface