diff --git a/dreambooth_gui.py b/dreambooth_gui.py index 393447c..7248c87 100644 --- a/dreambooth_gui.py +++ b/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, - caption_extention, + cache_latents, + 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, - caption_extention, + cache_latents, + 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 = 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, - ) - 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 = 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, - ], - ) + ( + button_open_config, + button_save_config, + button_save_as_config, + config_file_name, + ) = gradio_config() + + ( + pretrained_model_name_or_path, + v2, + v_parameterization, + save_model_as, + model_list, + ) = gradio_source_model() 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, - caption_extention, + cache_latents, + 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( diff --git a/finetune_gui.py b/finetune_gui.py index adae1d8..e0e918c 100644 --- a/finetune_gui.py +++ b/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,run_cmd_advanced_training, + gradio_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 += 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 += ' --enable_bucket' + run_cmd += f' --resolution={max_resolution}' + run_cmd += f' --min_bucket_reso={min_bucket_reso}' + run_cmd += f' --max_bucket_reso={max_bucket_reso}' + 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) - v_parameterization_input = gr.Checkbox( - label='v_parameterization', value=False - ) - 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, - ], - ) + ( + button_open_config, + button_save_config, + button_save_as_config, + config_file_name, + ) = gradio_config() + + ( + pretrained_model_name_or_path, + v2, + v_parameterization, + 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( - get_folder_path, outputs=output_dir_input - ) + output_dir_input_folder.click(get_folder_path, outputs=output_dir) - 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], - outputs=[train_dir_input], + inputs=[train_dir], + outputs=[train_dir], ) - image_folder_input.change( + image_folder.change( remove_doublequote, - inputs=[image_folder_input], - outputs=[image_folder_input], + inputs=[image_folder], + outputs=[image_folder], ) - output_dir_input.change( + output_dir.change( remove_doublequote, - inputs=[output_dir_input], - outputs=[output_dir_input], + inputs=[output_dir], + 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) - 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(): - 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( + dataset_repeats = gr.Textbox(label='Dataset repeats', value=40) + train_text_encoder = 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_changed, - # inputs=[color_aug], - # # outputs=[cache_latent], # 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 + color_aug.change( + color_aug_changed, + inputs=[color_aug], + outputs=[cache_latents], # Not applicable to fine_tune.py ) button_run = gr.Button('Train model') settings_list = [ - pretrained_model_name_or_path_input, - v2_input, - v_parameterization_input, - train_dir_input, - image_folder_input, - output_dir_input, - logging_dir_input, - max_resolution_input, + pretrained_model_name_or_path, + v2, + v_parameterization, + train_dir, + image_folder, + output_dir, + logging_dir, + 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, - lr_scheduler_input, - lr_warmup_input, - dataset_repeats_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, - train_text_encoder_input, + learning_rate, + lr_scheduler, + lr_warmup, + dataset_repeats, + train_batch_size, + epoch, + save_every_n_epochs, + mixed_precision, + save_precision, + seed, + num_cpu_threads_per_process, + train_text_encoder, create_caption, create_buckets, - save_model_as_dropdown, - caption_extention_input, + save_model_as, + 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) diff --git a/library/common_gui.py b/library/common_gui.py index a7b798b..7514854 100644 --- a/library/common_gui.py +++ b/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 + diff --git a/lora_gui.py b/lora_gui.py index 80b196e..eeba690 100644 --- a/lora_gui.py +++ b/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, - caption_extention, + cache_latents, + 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, - caption_extention, + cache_latents, + 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 = 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, - ) - # 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', - ) + ( + button_open_config, + button_save_config, + button_save_as_config, + config_file_name, + ) = gradio_config() - 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, - ], - ) + ( + pretrained_model_name_or_path, + v2, + v_parameterization, + save_model_as, + model_list, + ) = gradio_source_model() 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( - label='LR Scheduler', - choices=[ - 'constant', - 'constant_with_warmup', - 'cosine', - 'cosine_with_restarts', - 'linear', - 'polynomial', - ], - value='cosine', - ) - lr_warmup = gr.Textbox(label='LR warmup (% of steps)', value=10) + ( + 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(): 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, - caption_extention, + cache_latents, + 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,