Complete training code refactoring

dreambooth_gui.py

@@ -9,7 +9,6 @@ import math
 import os
 import subprocess
 import pathlib
-import shutil
 import argparse
 from library.common_gui import (
     get_folder_path,
@@ -19,9 +18,12 @@ from library.common_gui import (
     get_saveasfile_path,
     color_aug_changed,
     save_inference_file,
-    set_pretrained_model_name_or_path_input,
     gradio_advanced_training,
     run_cmd_advanced_training,
+    run_cmd_training,
+    gradio_training,
+    gradio_config,
+    gradio_source_model,
 )
 from library.dreambooth_folder_creation_gui import (
     gradio_dreambooth_folder_creation_tab,
@@ -56,8 +58,8 @@ def save_configuration(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
-    caption_extention,
+    caption_extension,
     enable_bucket,
     gradient_checkpointing,
     full_fp16,
@@ -77,8 +79,10 @@ def save_configuration(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,mem_eff_attn,
+    max_data_loader_n_workers,
+    mem_eff_attn,
     gradient_accumulation_steps,
+    model_list,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -138,8 +142,8 @@ def open_configuration(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
-    caption_extention,
+    caption_extension,
     enable_bucket,
     gradient_checkpointing,
     full_fp16,
@@ -159,8 +163,10 @@ def open_configuration(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,mem_eff_attn,
+    max_data_loader_n_workers,
+    mem_eff_attn,
     gradient_accumulation_steps,
+    model_list,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -172,7 +178,7 @@ def open_configuration(
         # load variables from JSON file
         with open(file_path, 'r') as f:
             my_data_db = json.load(f)
-            print("Loading config...")
+            print('Loading config...')
     else:
         file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
         my_data_db = {}
@@ -184,7 +190,7 @@ def open_configuration(
             values.append(my_data_db.get(key, value))
     return tuple(values)
 
-    
+
 def train_model(
     pretrained_model_name_or_path,
     v2,
@@ -204,7 +210,7 @@ def train_model(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
     caption_extension,
     enable_bucket,
     gradient_checkpointing,
@@ -225,8 +231,10 @@ def train_model(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,mem_eff_attn,
+    max_data_loader_n_workers,
+    mem_eff_attn,
     gradient_accumulation_steps,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
 ):
     if pretrained_model_name_or_path == '':
         msgbox('Source model information is missing')
@@ -321,8 +329,6 @@ def train_model(
         run_cmd += ' --v2'
     if v_parameterization:
         run_cmd += ' --v_parameterization'
-    if cache_latent:
-        run_cmd += ' --cache_latents'
     if enable_bucket:
         run_cmd += ' --enable_bucket'
     if no_token_padding:
@@ -339,18 +345,7 @@ def train_model(
         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' --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}"'
-    if not caption_extension == '':
-        run_cmd += f' --caption_extension={caption_extension}'
     if not stop_text_encoder_training == 0:
         run_cmd += (
             f' --stop_text_encoder_training={stop_text_encoder_training}'
@@ -365,14 +360,31 @@ def train_model(
         run_cmd += f' --vae="{vae}"'
     if not output_name == '':
         run_cmd += f' --output_name="{output_name}"'
-    if (int(max_token_length) > 75):
+    if int(max_token_length) > 75:
         run_cmd += f' --max_token_length={max_token_length}'
     if not max_train_epochs == '':
         run_cmd += f' --max_train_epochs="{max_train_epochs}"'
     if not max_data_loader_n_workers == '':
-        run_cmd += f' --max_data_loader_n_workers="{max_data_loader_n_workers}"'
+        run_cmd += (
+            f' --max_data_loader_n_workers="{max_data_loader_n_workers}"'
+        )
     if int(gradient_accumulation_steps) > 1:
         run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+    )
+
     run_cmd += run_cmd_advanced_training(
         max_train_epochs=max_train_epochs,
         max_data_loader_n_workers=max_data_loader_n_workers,
@@ -445,82 +457,20 @@ def dreambooth_tab(
     dummy_db_true = gr.Label(value=True, visible=False)
     dummy_db_false = gr.Label(value=False, visible=False)
     gr.Markdown('Train a custom model using kohya dreambooth python code...')
-    with gr.Accordion('Configuration file', open=False):
+    (
-        with gr.Row():
+        button_open_config,
-            button_open_config = gr.Button('Open 📂', elem_id='open_folder')
+        button_save_config,
-            button_save_config = gr.Button('Save 💾', elem_id='open_folder')
+        button_save_as_config,
-            button_save_as_config = gr.Button(
+        config_file_name,
-                'Save as... 💾', elem_id='open_folder'
+    ) = gradio_config()
-            )
+
-            config_file_name = gr.Textbox(
+    (
-                label='',
+        pretrained_model_name_or_path,
-                placeholder="type the configuration file path or use the 'Open' button above to select it...",
+        v2,
-                interactive=True,
+        v_parameterization,
-            )
+        save_model_as,
-    with gr.Tab('Source model'):
+        model_list,
-        # Define the input elements
+    ) = gradio_source_model()
-        with gr.Row():
-            pretrained_model_name_or_path = gr.Textbox(
-                label='Pretrained model name or path',
-                placeholder='enter the path to custom model or name of pretrained model',
-            )
-            pretrained_model_name_or_path_file = gr.Button(
-                document_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_file.click(
-                get_any_file_path,
-                inputs=[pretrained_model_name_or_path],
-                outputs=pretrained_model_name_or_path,
-            )
-            pretrained_model_name_or_path_folder = gr.Button(
-                folder_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_folder.click(
-                get_folder_path,
-                outputs=pretrained_model_name_or_path,
-            )
-            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',
-                ],
-            )
-            save_model_as = gr.Dropdown(
-                label='Save trained model as',
-                choices=[
-                    'same as source model',
-                    'ckpt',
-                    'diffusers',
-                    'diffusers_safetensors',
-                    'safetensors',
-                ],
-                value='same as source model',
-            )
-        with gr.Row():
-            v2 = gr.Checkbox(label='v2', value=True)
-            v_parameterization = gr.Checkbox(
-                label='v_parameterization', value=False
-            )
-        pretrained_model_name_or_path.change(
-            remove_doublequote,
-            inputs=[pretrained_model_name_or_path],
-            outputs=[pretrained_model_name_or_path],
-        )
-        model_list.change(
-            set_pretrained_model_name_or_path_input,
-            inputs=[model_list, v2, v_parameterization],
-            outputs=[
-                pretrained_model_name_or_path,
-                v2,
-                v_parameterization,
-            ],
-        )
 
     with gr.Tab('Folders'):
         with gr.Row():
@@ -591,71 +541,30 @@ def dreambooth_tab(
             outputs=[logging_dir],
         )
     with gr.Tab('Training parameters'):
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+        ) = gradio_training(
+            learning_rate_value='1e-5',
+            lr_scheduler_value='cosine',
+            lr_warmup_value='10',
+        )
         with gr.Row():
-            learning_rate = gr.Textbox(label='Learning rate', value=1e-6)
-            lr_scheduler = gr.Dropdown(
-                label='LR Scheduler',
-                choices=[
-                    'constant',
-                    'constant_with_warmup',
-                    'cosine',
-                    'cosine_with_restarts',
-                    'linear',
-                    'polynomial',
-                ],
-                value='constant',
-            )
-            lr_warmup = gr.Textbox(label='LR warmup', value=0)
-        with gr.Row():
-            train_batch_size = gr.Slider(
-                minimum=1,
-                maximum=32,
-                label='Train batch size',
-                value=1,
-                step=1,
-            )
-            epoch = gr.Textbox(label='Epoch', value=1)
-            save_every_n_epochs = gr.Textbox(
-                label='Save every N epochs', value=1
-            )
-        with gr.Row():
-            mixed_precision = gr.Dropdown(
-                label='Mixed precision',
-                choices=[
-                    'no',
-                    'fp16',
-                    'bf16',
-                ],
-                value='fp16',
-            )
-            save_precision = gr.Dropdown(
-                label='Save precision',
-                choices=[
-                    'float',
-                    'fp16',
-                    'bf16',
-                ],
-                value='fp16',
-            )
-            num_cpu_threads_per_process = 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 = gr.Textbox(label='Seed', value=1234)
             max_resolution = gr.Textbox(
                 label='Max resolution',
                 value='512,512',
                 placeholder='512,512',
             )
-        with gr.Row():
-            caption_extention = gr.Textbox(
-                label='Caption Extension',
-                placeholder='(Optional) Extension for caption files. default: .caption',
-            )
             stop_text_encoder_training = gr.Slider(
                 minimum=0,
                 maximum=100,
@@ -663,9 +572,7 @@ def dreambooth_tab(
                 step=1,
                 label='Stop text encoder training',
             )
-        with gr.Row():
             enable_bucket = gr.Checkbox(label='Enable buckets', value=True)
-            cache_latent = gr.Checkbox(label='Cache latent', value=True)
         with gr.Accordion('Advanced Configuration', open=False):
             with gr.Row():
                 no_token_padding = gr.Checkbox(
@@ -703,7 +610,7 @@ def dreambooth_tab(
             color_aug.change(
                 color_aug_changed,
                 inputs=[color_aug],
-                outputs=[cache_latent],
+                outputs=[cache_latents],
             )
     with gr.Tab('Tools'):
         gr.Markdown(
@@ -737,8 +644,8 @@ def dreambooth_tab(
         save_precision,
         seed,
         num_cpu_threads_per_process,
-        cache_latent,
+        cache_latents,
-        caption_extention,
+        caption_extension,
         enable_bucket,
         gradient_checkpointing,
         full_fp16,
@@ -758,8 +665,10 @@ def dreambooth_tab(
         output_name,
         max_token_length,
         max_train_epochs,
-        max_data_loader_n_workers,mem_eff_attn,
+        max_data_loader_n_workers,
+        mem_eff_attn,
         gradient_accumulation_steps,
+        model_list,
     ]
 
     button_open_config.click(

finetune_gui.py

@@ -4,17 +4,20 @@ import math
 import os
 import subprocess
 import pathlib
-import shutil
 import argparse
 from library.common_gui import (
     get_folder_path,
     get_file_path,
-    get_any_file_path,
     get_saveasfile_path,
     save_inference_file,
-    set_pretrained_model_name_or_path_input,
+    gradio_advanced_training,
-    gradio_advanced_training,run_cmd_advanced_training,
+    run_cmd_advanced_training,
+    gradio_training,
+    run_cmd_advanced_training,
+    gradio_config,
+    gradio_source_model,
     color_aug_changed,
+    run_cmd_training,
 )
 from library.utilities import utilities_tab
 
@@ -70,11 +73,16 @@ def save_configuration(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,full_fp16,color_aug,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,
+    cache_latents,
+    use_latent_files,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
-    
+
     original_file_path = file_path
 
     save_as_bool = True if save_as.get('label') == 'True' else False
@@ -155,11 +163,16 @@ def open_config_file(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,full_fp16,color_aug,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,
+    cache_latents,
+    use_latent_files,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
-    
+
     original_file_path = file_path
     file_path = get_file_path(file_path)
 
@@ -171,7 +184,7 @@ def open_config_file(
     else:
         file_path = original_file_path   # In case a file_path was provided and the user decide to cancel the open action
         my_data_ft = {}
-    
+
     values = [file_path]
     for key, value in parameters:
         # Set the value in the dictionary to the corresponding value in `my_data_ft`, or the default value if not found
@@ -225,7 +238,12 @@ def train_model(
     output_name,
     max_token_length,
     max_train_epochs,
-    max_data_loader_n_workers,full_fp16,color_aug,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
+    cache_latents,
+    use_latent_files,
 ):
     # create caption json file
     if generate_caption_database:
@@ -236,7 +254,7 @@ def train_model(
             f'./venv/Scripts/python.exe finetune/merge_captions_to_metadata.py'
         )
         if caption_extension == '':
-            run_cmd += f' --caption_extension=".txt"'
+            run_cmd += f' --caption_extension=".caption"'
         else:
             run_cmd += f' --caption_extension={caption_extension}'
         run_cmd += f' "{image_folder}"'
@@ -305,21 +323,22 @@ def train_model(
     run_cmd += (
         f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
     )
-    run_cmd += f' --in_json="{train_dir}/{latent_metadata_filename}"'
+    if use_latent_files == 'Yes':
+        run_cmd += f' --in_json="{train_dir}/{latent_metadata_filename}"'
+    else:
+        run_cmd += f' --in_json="{train_dir}/{caption_metadata_filename}"'
     run_cmd += f' --train_data_dir="{image_folder}"'
     run_cmd += f' --output_dir="{output_dir}"'
     if not logging_dir == '':
         run_cmd += f' --logging_dir="{logging_dir}"'
-    run_cmd += f' --train_batch_size={train_batch_size}'
     run_cmd += f' --dataset_repeats={dataset_repeats}'
     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 += ' --enable_bucket'
-    run_cmd += f' --max_train_steps={max_train_steps}'
+    run_cmd += f' --resolution={max_resolution}'
-    run_cmd += f' --mixed_precision={mixed_precision}'
+    run_cmd += f' --min_bucket_reso={min_bucket_reso}'
-    run_cmd += f' --save_every_n_epochs={save_every_n_epochs}'
+    run_cmd += f' --max_bucket_reso={max_bucket_reso}'
-    run_cmd += f' --seed={seed}'
+
-    run_cmd += f' --save_precision={save_precision}'
     if not save_model_as == 'same as source model':
         run_cmd += f' --save_model_as={save_model_as}'
     if int(gradient_accumulation_steps) > 1:
@@ -330,8 +349,23 @@ def train_model(
     #     run_cmd += f' --resume={resume}'
     if not output_name == '':
         run_cmd += f' --output_name="{output_name}"'
-    if (int(max_token_length) > 75):
+    if int(max_token_length) > 75:
         run_cmd += f' --max_token_length={max_token_length}'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+    )
+
     run_cmd += run_cmd_advanced_training(
         max_train_epochs=max_train_epochs,
         max_data_loader_n_workers=max_data_loader_n_workers,
@@ -396,99 +430,34 @@ def finetune_tab():
     dummy_ft_true = gr.Label(value=True, visible=False)
     dummy_ft_false = gr.Label(value=False, visible=False)
     gr.Markdown('Train a custom model using kohya finetune python code...')
-    with gr.Accordion('Configuration file', open=False):
-        with gr.Row():
-            button_open_config = gr.Button(
-                f'Open {folder_symbol}', elem_id='open_folder'
-            )
-            button_save_config = gr.Button(
-                f'Save {save_style_symbol}', elem_id='open_folder'
-            )
-            button_save_as_config = gr.Button(
-                f'Save as... {save_style_symbol}',
-                elem_id='open_folder',
-            )
-        config_file_name = gr.Textbox(
-            label='', placeholder='type 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',
-            )
-            pretrained_model_name_or_path_file = gr.Button(
-                document_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_file.click(
-                get_any_file_path,
-                inputs=pretrained_model_name_or_path_input,
-                outputs=pretrained_model_name_or_path_input,
-            )
-            pretrained_model_name_or_path_folder = gr.Button(
-                folder_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_folder.click(
-                get_folder_path,
-                inputs=pretrained_model_name_or_path_input,
-                outputs=pretrained_model_name_or_path_input,
-            )
-            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',
-                ],
-            )
-            save_model_as_dropdown = gr.Dropdown(
-                label='Save trained model as',
-                choices=[
-                    'same as source model',
-                    'ckpt',
-                    'diffusers',
-                    'diffusers_safetensors',
-                    'safetensors',
-                ],
-                value='same as source model',
-            )
 
-        with gr.Row():
+    (
-            v2_input = gr.Checkbox(label='v2', value=True)
+        button_open_config,
-            v_parameterization_input = gr.Checkbox(
+        button_save_config,
-                label='v_parameterization', value=False
+        button_save_as_config,
-            )
+        config_file_name,
-        model_list.change(
+    ) = gradio_config()
-            set_pretrained_model_name_or_path_input,
+
-            inputs=[model_list, v2_input, v_parameterization_input],
+    (
-            outputs=[
+        pretrained_model_name_or_path,
-                pretrained_model_name_or_path_input,
+        v2,
-                v2_input,
+        v_parameterization,
-                v_parameterization_input,
+        save_model_as,
-            ],
+        model_list,
-        )
+    ) = gradio_source_model()
+
     with gr.Tab('Folders'):
         with gr.Row():
-            train_dir_input = gr.Textbox(
+            train_dir = gr.Textbox(
                 label='Training config folder',
                 placeholder='folder where the training configuration files will be saved',
             )
             train_dir_folder = gr.Button(
                 folder_symbol, elem_id='open_folder_small'
             )
-            train_dir_folder.click(get_folder_path, outputs=train_dir_input)
+            train_dir_folder.click(get_folder_path, outputs=train_dir)
 
-            image_folder_input = gr.Textbox(
+            image_folder = gr.Textbox(
                 label='Training Image folder',
                 placeholder='folder where the training images are located',
             )
@@ -496,21 +465,19 @@ def finetune_tab():
                 folder_symbol, elem_id='open_folder_small'
             )
             image_folder_input_folder.click(
-                get_folder_path, outputs=image_folder_input
+                get_folder_path, outputs=image_folder
             )
         with gr.Row():
-            output_dir_input = gr.Textbox(
+            output_dir = gr.Textbox(
                 label='Model output folder',
                 placeholder='folder where the model will be saved',
             )
             output_dir_input_folder = gr.Button(
                 folder_symbol, elem_id='open_folder_small'
             )
-            output_dir_input_folder.click(
+            output_dir_input_folder.click(get_folder_path, outputs=output_dir)
-                get_folder_path, outputs=output_dir_input
-            )
 
-            logging_dir_input = gr.Textbox(
+            logging_dir = gr.Textbox(
                 label='Logging folder',
                 placeholder='Optional: enable logging and output TensorBoard log to this folder',
             )
@@ -518,7 +485,7 @@ def finetune_tab():
                 folder_symbol, elem_id='open_folder_small'
             )
             logging_dir_input_folder.click(
-                get_folder_path, outputs=logging_dir_input
+                get_folder_path, outputs=logging_dir
             )
         with gr.Row():
             output_name = gr.Textbox(
@@ -527,24 +494,24 @@ def finetune_tab():
                 value='last',
                 interactive=True,
             )
-        train_dir_input.change(
+        train_dir.change(
             remove_doublequote,
-            inputs=[train_dir_input],
+            inputs=[train_dir],
-            outputs=[train_dir_input],
+            outputs=[train_dir],
         )
-        image_folder_input.change(
+        image_folder.change(
             remove_doublequote,
-            inputs=[image_folder_input],
+            inputs=[image_folder],
-            outputs=[image_folder_input],
+            outputs=[image_folder],
         )
-        output_dir_input.change(
+        output_dir.change(
             remove_doublequote,
-            inputs=[output_dir_input],
+            inputs=[output_dir],
-            outputs=[output_dir_input],
+            outputs=[output_dir],
         )
     with gr.Tab('Dataset preparation'):
         with gr.Row():
-            max_resolution_input = gr.Textbox(
+            max_resolution = gr.Textbox(
                 label='Resolution (width,height)', value='512,512'
             )
             min_bucket_reso = gr.Textbox(
@@ -554,6 +521,21 @@ def finetune_tab():
                 label='Max bucket resolution', value='1024'
             )
             batch_size = gr.Textbox(label='Batch size', value='1')
+        with gr.Row():
+            create_caption = gr.Checkbox(
+                label='Generate caption metadata', value=True
+            )
+            create_buckets = gr.Checkbox(
+                label='Generate image buckets metadata', value=True
+            )
+            use_latent_files = gr.Dropdown(
+                label='Use latent files',
+                choices=[
+                    'No',
+                    'Yes',
+                ],
+                value='Yes',
+            )
         with gr.Accordion('Advanced parameters', open=False):
             with gr.Row():
                 caption_metadata_filename = gr.Textbox(
@@ -564,69 +546,23 @@ def finetune_tab():
                 )
                 full_path = gr.Checkbox(label='Use full path', value=True)
     with gr.Tab('Training parameters'):
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+        ) = gradio_training(learning_rate_value='1e-5')
         with gr.Row():
-            learning_rate_input = gr.Textbox(label='Learning rate', value=1e-6)
+            dataset_repeats = gr.Textbox(label='Dataset repeats', value=40)
-            lr_scheduler_input = gr.Dropdown(
+            train_text_encoder = gr.Checkbox(
-                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():
-            dataset_repeats_input = gr.Textbox(
-                label='Dataset repeats', value=40
-            )
-            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(),
-            )
-            seed_input = gr.Textbox(label='Seed', value=1234)
-        with gr.Row():
-            caption_extention_input = gr.Textbox(
-                label='Caption Extension',
-                placeholder='(Optional) Extension for caption files. default: .txt',
-            )
-            train_text_encoder_input = gr.Checkbox(
                 label='Train text encoder', value=True
             )
         with gr.Accordion('Advanced parameters', open=False):
@@ -650,31 +586,23 @@ def finetune_tab():
                 max_train_epochs,
                 max_data_loader_n_workers,
             ) = gradio_advanced_training()
-            # color_aug.change(
+            color_aug.change(
-            #     color_aug_changed,
+                color_aug_changed,
-            #     inputs=[color_aug],
+                inputs=[color_aug],
-            #     # outputs=[cache_latent], # Not applicable to fine_tune.py
+                outputs=[cache_latents],  # Not applicable to fine_tune.py
-            # )
-    with gr.Box():
-        with gr.Row():
-            create_caption = gr.Checkbox(
-                label='Generate caption metadata', value=True
-            )
-            create_buckets = gr.Checkbox(
-                label='Generate image buckets metadata', value=True
             )
 
     button_run = gr.Button('Train model')
 
     settings_list = [
-        pretrained_model_name_or_path_input,
+        pretrained_model_name_or_path,
-        v2_input,
+        v2,
-        v_parameterization_input,
+        v_parameterization,
-        train_dir_input,
+        train_dir,
-        image_folder_input,
+        image_folder,
-        output_dir_input,
+        output_dir,
-        logging_dir_input,
+        logging_dir,
-        max_resolution_input,
+        max_resolution,
         min_bucket_reso,
         max_bucket_reso,
         batch_size,
@@ -682,22 +610,22 @@ def finetune_tab():
         caption_metadata_filename,
         latent_metadata_filename,
         full_path,
-        learning_rate_input,
+        learning_rate,
-        lr_scheduler_input,
+        lr_scheduler,
-        lr_warmup_input,
+        lr_warmup,
-        dataset_repeats_input,
+        dataset_repeats,
-        train_batch_size_input,
+        train_batch_size,
-        epoch_input,
+        epoch,
-        save_every_n_epochs_input,
+        save_every_n_epochs,
-        mixed_precision_input,
+        mixed_precision,
-        save_precision_input,
+        save_precision,
-        seed_input,
+        seed,
-        num_cpu_threads_per_process_input,
+        num_cpu_threads_per_process,
-        train_text_encoder_input,
+        train_text_encoder,
         create_caption,
         create_buckets,
-        save_model_as_dropdown,
+        save_model_as,
-        caption_extention_input,
+        caption_extension,
         use_8bit_adam,
         xformers,
         clip_skip,
@@ -710,7 +638,12 @@ def finetune_tab():
         output_name,
         max_token_length,
         max_train_epochs,
-        max_data_loader_n_workers,full_fp16,color_aug,
+        max_data_loader_n_workers,
+        full_fp16,
+        color_aug,
+        model_list,
+        cache_latents,
+        use_latent_files,
     ]
 
     button_run.click(train_model, inputs=settings_list)

library/common_gui.py

@@ -4,6 +4,10 @@ import gradio as gr
 from easygui import msgbox
 import shutil
 
+folder_symbol = '\U0001f4c2'  # 📂
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+document_symbol = '\U0001F4C4'   # 📄
 
 def get_dir_and_file(file_path):
     dir_path, file_name = os.path.split(file_path)
@@ -300,6 +304,208 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
     ###
     ### Gradio common GUI section
     ###
+    
+def gradio_config():
+    with gr.Accordion('Configuration file', 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 the configuration file path or use the 'Open' button above to select it...",
+                interactive=True,
+            )
+    return (button_open_config, button_save_config, button_save_as_config, config_file_name)
+
+def gradio_source_model():
+    with gr.Tab('Source model'):
+        # Define the input elements
+        with gr.Row():
+            pretrained_model_name_or_path = gr.Textbox(
+                label='Pretrained model name or path',
+                placeholder='enter the path to custom model or name of pretrained model',
+            )
+            pretrained_model_name_or_path_file = gr.Button(
+                document_symbol, elem_id='open_folder_small'
+            )
+            pretrained_model_name_or_path_file.click(
+                get_any_file_path,
+                inputs=pretrained_model_name_or_path,
+                outputs=pretrained_model_name_or_path,
+            )
+            pretrained_model_name_or_path_folder = gr.Button(
+                folder_symbol, elem_id='open_folder_small'
+            )
+            pretrained_model_name_or_path_folder.click(
+                get_folder_path,
+                inputs=pretrained_model_name_or_path,
+                outputs=pretrained_model_name_or_path,
+            )
+            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',
+                ],
+            )
+            save_model_as = gr.Dropdown(
+                label='Save trained model as',
+                choices=[
+                    'same as source model',
+                    'ckpt',
+                    'diffusers',
+                    'diffusers_safetensors',
+                    'safetensors',
+                ],
+                value='same as source model',
+            )
+
+        with gr.Row():
+            v2 = gr.Checkbox(label='v2', value=True)
+            v_parameterization = gr.Checkbox(
+                label='v_parameterization', value=False
+            )
+        model_list.change(
+            set_pretrained_model_name_or_path_input,
+            inputs=[model_list, v2, v_parameterization],
+            outputs=[
+                pretrained_model_name_or_path,
+                v2,
+                v_parameterization,
+            ],
+        )
+    return (pretrained_model_name_or_path, v2, v_parameterization, save_model_as, model_list)
+
+def gradio_training(learning_rate_value='1e-6', lr_scheduler_value='constant', lr_warmup_value='0'):
+    with gr.Row():
+        train_batch_size = gr.Slider(
+            minimum=1,
+            maximum=32,
+            label='Train batch size',
+            value=1,
+            step=1,
+        )
+        epoch = gr.Textbox(label='Epoch', value=1)
+        save_every_n_epochs = gr.Textbox(
+            label='Save every N epochs', value=1
+        )
+        caption_extension = gr.Textbox(
+            label='Caption Extension',
+            placeholder='(Optional) Extension for caption files. default: .caption',
+        )
+    with gr.Row():
+        mixed_precision = gr.Dropdown(
+            label='Mixed precision',
+            choices=[
+                'no',
+                'fp16',
+                'bf16',
+            ],
+            value='fp16',
+        )
+        save_precision = gr.Dropdown(
+            label='Save precision',
+            choices=[
+                'float',
+                'fp16',
+                'bf16',
+            ],
+            value='fp16',
+        )
+        num_cpu_threads_per_process = gr.Slider(
+            minimum=1,
+            maximum=os.cpu_count(),
+            step=1,
+            label='Number of CPU threads per process',
+            value=os.cpu_count(),
+        )
+        seed = gr.Textbox(label='Seed', value=1234)
+    with gr.Row():
+        learning_rate = gr.Textbox(label='Learning rate', value=learning_rate_value)
+        lr_scheduler = gr.Dropdown(
+            label='LR Scheduler',
+            choices=[
+                'constant',
+                'constant_with_warmup',
+                'cosine',
+                'cosine_with_restarts',
+                'linear',
+                'polynomial',
+            ],
+            value=lr_scheduler_value,
+        )
+        lr_warmup = gr.Textbox(label='LR warmup (% of steps)', value=lr_warmup_value)
+        cache_latents = gr.Checkbox(label='Cache latent', value=True)
+    return (
+        learning_rate,
+        lr_scheduler,
+        lr_warmup,
+        train_batch_size,
+        epoch,
+        save_every_n_epochs,
+        mixed_precision,
+        save_precision,
+        num_cpu_threads_per_process,
+        seed,
+        caption_extension,
+        cache_latents,
+    )
+
+def run_cmd_training(**kwargs):
+    options = [
+        f' --learning_rate="{kwargs.get("learning_rate", "")}"'
+        if kwargs.get('learning_rate')
+        else '',
+        
+        f' --lr_scheduler="{kwargs.get("lr_scheduler", "")}"'
+        if kwargs.get('lr_scheduler')
+        else '',
+        
+        f' --lr_warmup_steps="{kwargs.get("lr_warmup_steps", "")}"'
+        if kwargs.get('lr_warmup_steps')
+        else '',
+        
+        f' --train_batch_size="{kwargs.get("train_batch_size", "")}"'
+        if kwargs.get('train_batch_size')
+        else '',
+        
+        f' --max_train_steps="{kwargs.get("max_train_steps", "")}"'
+        if kwargs.get('max_train_steps')
+        else '',
+        
+        f' --save_every_n_epochs="{kwargs.get("save_every_n_epochs", "")}"'
+        if kwargs.get('save_every_n_epochs')
+        else '',
+        
+        f' --mixed_precision="{kwargs.get("mixed_precision", "")}"'
+        if kwargs.get('mixed_precision')
+        else '',
+        
+        f' --save_precision="{kwargs.get("save_precision", "")}"'
+        if kwargs.get('save_precision')
+        else '',
+        
+        f' --seed="{kwargs.get("seed", "")}"'
+        if kwargs.get('seed')
+        else '',
+        
+        f' --caption_extension="{kwargs.get("caption_extension", "")}"'
+        if kwargs.get('caption_extension')
+        else '',
+        
+        ' --cache_latents' if kwargs.get('cache_latents') else '',
+        
+    ]
+    run_cmd = ''.join(options)
+    return run_cmd
 
 
 def gradio_advanced_training():
@@ -368,7 +574,6 @@ def gradio_advanced_training():
         max_data_loader_n_workers,
     )
 
-
 def run_cmd_advanced_training(**kwargs):
     options = [
         f' --max_train_epochs="{kwargs.get("max_train_epochs", "")}"'
@@ -412,3 +617,4 @@ def run_cmd_advanced_training(**kwargs):
     ]
     run_cmd = ''.join(options)
     return run_cmd
+

lora_gui.py

@@ -9,7 +9,6 @@ import math
 import os
 import subprocess
 import pathlib
-import shutil
 import argparse
 from library.common_gui import (
     get_folder_path,
@@ -19,9 +18,12 @@ from library.common_gui import (
     get_saveasfile_path,
     color_aug_changed,
     save_inference_file,
-    set_pretrained_model_name_or_path_input,
     gradio_advanced_training,
     run_cmd_advanced_training,
+    gradio_training,
+    gradio_config,
+    gradio_source_model,
+    run_cmd_training,
 )
 from library.dreambooth_folder_creation_gui import (
     gradio_dreambooth_folder_creation_tab,
@@ -48,6 +50,7 @@ def save_configuration(
     reg_data_dir,
     output_dir,
     max_resolution,
+    learning_rate,
     lr_scheduler,
     lr_warmup,
     train_batch_size,
@@ -57,8 +60,8 @@ def save_configuration(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
-    caption_extention,
+    caption_extension,
     enable_bucket,
     gradient_checkpointing,
     full_fp16,
@@ -134,6 +137,7 @@ def open_configuration(
     reg_data_dir,
     output_dir,
     max_resolution,
+    learning_rate,
     lr_scheduler,
     lr_warmup,
     train_batch_size,
@@ -143,8 +147,8 @@ def open_configuration(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
-    caption_extention,
+    caption_extension,
     enable_bucket,
     gradient_checkpointing,
     full_fp16,
@@ -204,6 +208,7 @@ def train_model(
     reg_data_dir,
     output_dir,
     max_resolution,
+    learning_rate,
     lr_scheduler,
     lr_warmup,
     train_batch_size,
@@ -213,7 +218,7 @@ def train_model(
     save_precision,
     seed,
     num_cpu_threads_per_process,
-    cache_latent,
+    cache_latents,
     caption_extension,
     enable_bucket,
     gradient_checkpointing,
@@ -336,8 +341,6 @@ def train_model(
         run_cmd += ' --v2'
     if v_parameterization:
         run_cmd += ' --v_parameterization'
-    if cache_latent:
-        run_cmd += ' --cache_latents'
     if enable_bucket:
         run_cmd += ' --enable_bucket'
     if no_token_padding:
@@ -350,28 +353,15 @@ def train_model(
         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}"'
-    if not caption_extension == '':
-        run_cmd += f' --caption_extension={caption_extension}'
     if not stop_text_encoder_training == 0:
         run_cmd += (
             f' --stop_text_encoder_training={stop_text_encoder_training}'
         )
     if not save_model_as == 'same as source model':
         run_cmd += f' --save_model_as={save_model_as}'
-    # if not resume == '':
-    #     run_cmd += f' --resume="{resume}"'
     if not float(prior_loss_weight) == 1.0:
         run_cmd += f' --prior_loss_weight={prior_loss_weight}'
     run_cmd += f' --network_module=networks.lora'
@@ -383,21 +373,28 @@ def train_model(
         run_cmd += f' --unet_lr={unet_lr}'
     else:
         run_cmd += f' --network_train_text_encoder_only'
-    # if network_train == 'Text encoder only':
-    #     run_cmd += f' --network_train_text_encoder_only'
-    # elif network_train == 'Unet only':
-    #     run_cmd += f' --network_train_unet_only'
     run_cmd += f' --network_dim={network_dim}'
     if not lora_network_weights == '':
         run_cmd += f' --network_weights="{lora_network_weights}"'
     if int(gradient_accumulation_steps) > 1:
         run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
-    # if not vae == '':
-    #     run_cmd += f' --vae="{vae}"'
     if not output_name == '':
         run_cmd += f' --output_name="{output_name}"'
-    # if (int(max_token_length) > 75):
+
-    #     run_cmd += f' --max_token_length={max_token_length}'
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+    )
+
     run_cmd += run_cmd_advanced_training(
         max_train_epochs=max_train_epochs,
         max_data_loader_n_workers=max_data_loader_n_workers,
@@ -472,88 +469,20 @@ def lora_tab(
     gr.Markdown(
         'Train a custom model using kohya train network LoRA python code...'
     )
-    with gr.Accordion('Configuration file', open=False):
+    (
-        with gr.Row():
+        button_open_config,
-            button_open_config = gr.Button('Open 📂', elem_id='open_folder')
+        button_save_config,
-            button_save_config = gr.Button('Save 💾', elem_id='open_folder')
+        button_save_as_config,
-            button_save_as_config = gr.Button(
+        config_file_name,
-                'Save as... 💾', elem_id='open_folder'
+    ) = gradio_config()
-            )
-            config_file_name = gr.Textbox(
-                label='',
-                placeholder="type the configuration file path or use the 'Open' button above to select it...",
-                interactive=True,
-            )
-        # 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 = gr.Textbox(
-                label='Pretrained model name or path',
-                placeholder='enter the path to custom model or name of pretrained model',
-            )
-            pretrained_model_name_or_path_file = gr.Button(
-                document_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_file.click(
-                get_any_file_path,
-                inputs=[pretrained_model_name_or_path],
-                outputs=pretrained_model_name_or_path,
-            )
-            pretrained_model_name_or_path_folder = gr.Button(
-                folder_symbol, elem_id='open_folder_small'
-            )
-            pretrained_model_name_or_path_folder.click(
-                get_folder_path,
-                outputs=pretrained_model_name_or_path,
-            )
-            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',
-                ],
-            )
-            save_model_as_dropdown = gr.Dropdown(
-                label='Save trained model as',
-                choices=[
-                    'same as source model',
-                    'ckpt',
-                    'diffusers',
-                    'diffusers_safetensors',
-                    'safetensors',
-                ],
-                value='same as source model',
-            )
 
-        with gr.Row():
+    (
-            v2 = gr.Checkbox(label='v2', value=True)
+        pretrained_model_name_or_path,
-            v_parameterization = gr.Checkbox(
+        v2,
-                label='v_parameterization', value=False
+        v_parameterization,
-            )
+        save_model_as,
-        pretrained_model_name_or_path.change(
+        model_list,
-            remove_doublequote,
+    ) = gradio_source_model()
-            inputs=[pretrained_model_name_or_path],
-            outputs=[pretrained_model_name_or_path],
-        )
-        model_list.change(
-            set_pretrained_model_name_or_path_input,
-            inputs=[model_list, v2, v_parameterization],
-            outputs=[
-                pretrained_model_name_or_path,
-                v2,
-                v_parameterization,
-            ],
-        )
 
     with gr.Tab('Folders'):
         with gr.Row():
@@ -625,20 +554,24 @@ def lora_tab(
                 inputs=[lora_network_weights],
                 outputs=lora_network_weights,
             )
-        with gr.Row():
+        (
-            lr_scheduler = gr.Dropdown(
+            learning_rate,
-                label='LR Scheduler',
+            lr_scheduler,
-                choices=[
+            lr_warmup,
-                    'constant',
+            train_batch_size,
-                    'constant_with_warmup',
+            epoch,
-                    'cosine',
+            save_every_n_epochs,
-                    'cosine_with_restarts',
+            mixed_precision,
-                    'linear',
+            save_precision,
-                    'polynomial',
+            num_cpu_threads_per_process,
-                ],
+            seed,
-                value='cosine',
+            caption_extension,
-            )
+            cache_latents,
-            lr_warmup = gr.Textbox(label='LR warmup (% of steps)', value=10)
+        ) = gradio_training(
+            learning_rate_value='1e-5',
+            lr_scheduler_value='cosine',
+            lr_warmup_value='10',
+        )
         with gr.Row():
             text_encoder_lr = gr.Textbox(
                 label='Text Encoder learning rate',
@@ -659,55 +592,11 @@ def lora_tab(
                 interactive=True,
             )
         with gr.Row():
-            train_batch_size = gr.Slider(
-                minimum=1,
-                maximum=32,
-                label='Train batch size',
-                value=1,
-                step=1,
-            )
-            epoch = gr.Textbox(label='Epoch', value=1)
-            save_every_n_epochs = gr.Textbox(
-                label='Save every N epochs', value=1
-            )
-        with gr.Row():
-            mixed_precision = gr.Dropdown(
-                label='Mixed precision',
-                choices=[
-                    'no',
-                    'fp16',
-                    'bf16',
-                ],
-                value='fp16',
-            )
-            save_precision = gr.Dropdown(
-                label='Save precision',
-                choices=[
-                    'float',
-                    'fp16',
-                    'bf16',
-                ],
-                value='fp16',
-            )
-            num_cpu_threads_per_process = 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 = gr.Textbox(label='Seed', value=1234)
             max_resolution = gr.Textbox(
                 label='Max resolution',
                 value='512,512',
                 placeholder='512,512',
             )
-        with gr.Row():
-            caption_extention = gr.Textbox(
-                label='Caption Extension',
-                placeholder='(Optional) Extension for caption files. default: .caption',
-            )
             stop_text_encoder_training = gr.Slider(
                 minimum=0,
                 maximum=100,
@@ -715,9 +604,7 @@ def lora_tab(
                 step=1,
                 label='Stop text encoder training',
             )
-        with gr.Row():
             enable_bucket = gr.Checkbox(label='Enable buckets', value=True)
-            cache_latent = gr.Checkbox(label='Cache latent', value=True)
         with gr.Accordion('Advanced Configuration', open=False):
             with gr.Row():
                 no_token_padding = gr.Checkbox(
@@ -749,7 +636,7 @@ def lora_tab(
             color_aug.change(
                 color_aug_changed,
                 inputs=[color_aug],
-                outputs=[cache_latent],
+                outputs=[cache_latents],
             )
 
     with gr.Tab('Tools'):
@@ -776,6 +663,7 @@ def lora_tab(
         reg_data_dir,
         output_dir,
         max_resolution,
+        learning_rate,
         lr_scheduler,
         lr_warmup,
         train_batch_size,
@@ -785,8 +673,8 @@ def lora_tab(
         save_precision,
         seed,
         num_cpu_threads_per_process,
-        cache_latent,
+        cache_latents,
-        caption_extention,
+        caption_extension,
         enable_bucket,
         gradient_checkpointing,
         full_fp16,
@@ -794,7 +682,7 @@ def lora_tab(
         stop_text_encoder_training,
         use_8bit_adam,
         xformers,
-        save_model_as_dropdown,
+        save_model_as,
         shuffle_caption,
         save_state,
         resume,