From a4262c0a66edd5cfb54b554f84098254da779b89 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Sun, 8 Jan 2023 20:55:41 -0500 Subject: [PATCH] - Add vae support to dreambooth GUI - Add gradient_checkpointing, gradient_accumulation_steps, mem_eff_attn, shuffle_caption to finetune GUI - Add gradient_accumulation_steps, mem_eff_attn to dreambooth lora gui --- README.md | 4 ++++ dreambooth_gui.py | 14 +++++++++++++ finetune_gui.py | 44 +++++++++++++++++++++++++++++++++++++++++ lora_gui.py | 50 ++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 5ed7ea9..e9ba1ad 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,10 @@ Once you have created the LoRA network you can generate images via auto1111 by i ## Change history +* 2023/01/09 (v19.4.3): + - Add vae support to dreambooth GUI + - Add gradient_checkpointing, gradient_accumulation_steps, mem_eff_attn, shuffle_caption to finetune GUI + - Add gradient_accumulation_steps, mem_eff_attn to dreambooth lora gui * 2023/01/08 (v19.4.2): - Add find/replace option to Basic Caption utility - Add resume training and save_state option to finetune UI diff --git a/dreambooth_gui.py b/dreambooth_gui.py index e2166a7..75e6868 100644 --- a/dreambooth_gui.py +++ b/dreambooth_gui.py @@ -69,6 +69,7 @@ def save_configuration( color_aug, flip_aug, clip_skip, + vae, ): original_file_path = file_path @@ -124,6 +125,7 @@ def save_configuration( 'color_aug': color_aug, 'flip_aug': flip_aug, 'clip_skip': clip_skip, + 'vae': vae, } # Save the data to the selected file @@ -170,6 +172,7 @@ def open_configuration( color_aug, flip_aug, clip_skip, + vae, ): original_file_path = file_path @@ -226,6 +229,7 @@ def open_configuration( my_data.get('color_aug', color_aug), my_data.get('flip_aug', flip_aug), my_data.get('clip_skip', clip_skip), + my_data.get('vae', vae), ) @@ -265,6 +269,7 @@ def train_model( color_aug, flip_aug, clip_skip, + vae, ): def save_inference_file(output_dir, v2, v_parameterization): # Copy inference model for v2 if required @@ -430,6 +435,8 @@ def train_model( run_cmd += f' --prior_loss_weight={prior_loss_weight}' if int(clip_skip) > 1: run_cmd += f' --clip_skip={str(clip_skip)}' + if not vae == '': + run_cmd += f' --vae="{vae}"' print(run_cmd) # Run the command @@ -806,6 +813,12 @@ def dreambooth_tab( prior_loss_weight = gr.Number( label='Prior loss weight', value=1.0 ) + vae = gr.Textbox( + label='VAE', + placeholder='(Optiona) path to checkpoint of vae to replace for training', + ) + vae_button = gr.Button('📂', elem_id='open_folder_small') + vae_button.click(get_any_file_path, outputs=vae) with gr.Tab('Tools'): gr.Markdown( 'This section provide Dreambooth tools to help setup your dataset...' @@ -855,6 +868,7 @@ def dreambooth_tab( color_aug, flip_aug, clip_skip, + vae, ] button_open_config.click( diff --git a/finetune_gui.py b/finetune_gui.py index 76b7298..cd1f1ff 100644 --- a/finetune_gui.py +++ b/finetune_gui.py @@ -59,6 +59,10 @@ def save_configuration( clip_skip, save_state, resume, + gradient_checkpointing, + gradient_accumulation_steps, + mem_eff_attn, + shuffle_caption, ): original_file_path = file_path @@ -115,6 +119,10 @@ def save_configuration( 'clip_skip': clip_skip, 'save_state': save_state, 'resume': resume, + 'gradient_checkpointing': gradient_checkpointing, + 'gradient_accumulation_steps': gradient_accumulation_steps, + 'mem_eff_attn': mem_eff_attn, + 'shuffle_caption': shuffle_caption, } # Save the data to the selected file @@ -162,6 +170,10 @@ def open_config_file( clip_skip, save_state, resume, + gradient_checkpointing, + gradient_accumulation_steps, + mem_eff_attn, + shuffle_caption, ): original_file_path = file_path file_path = get_file_path(file_path) @@ -218,6 +230,10 @@ def open_config_file( my_data.get('clip_skip', clip_skip), my_data.get('save_state', save_state), my_data.get('resume', resume), + my_data.get('gradient_checkpointing', gradient_checkpointing), + my_data.get('gradient_accumulation_steps', gradient_accumulation_steps), + my_data.get('mem_eff_attn', mem_eff_attn), + my_data.get('shuffle_caption', shuffle_caption), ) @@ -258,6 +274,10 @@ def train_model( clip_skip, save_state, resume, + gradient_checkpointing, + gradient_accumulation_steps, + mem_eff_attn, + shuffle_caption, ): def save_inference_file(output_dir, v2, v_parameterization): # Copy inference model for v2 if required @@ -353,6 +373,12 @@ def train_model( run_cmd += f' --use_8bit_adam' if xformers: run_cmd += f' --xformers' + if gradient_checkpointing: + run_cmd += ' --gradient_checkpointing' + if mem_eff_attn: + run_cmd += ' --mem_eff_attn' + if shuffle_caption: + run_cmd += ' --shuffle_caption' run_cmd += ( f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"' ) @@ -375,6 +401,8 @@ def train_model( run_cmd += f' --save_model_as={save_model_as}' if int(clip_skip) > 1: run_cmd += f' --clip_skip={str(clip_skip)}' + if int(gradient_accumulation_steps) > 1: + run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}' if save_state: run_cmd += ' --save_state' if not resume == '': @@ -712,6 +740,12 @@ def finetune_tab(): clip_skip = gr.Slider( label='Clip skip', value='1', minimum=1, maximum=12, step=1 ) + mem_eff_attn = gr.Checkbox( + label='Memory efficient attention', value=False + ) + shuffle_caption = gr.Checkbox( + label='Shuffle caption', value=False + ) with gr.Row(): save_state = gr.Checkbox( label='Save training state', value=False @@ -722,6 +756,12 @@ def finetune_tab(): ) resume_button = gr.Button('📂', elem_id='open_folder_small') resume_button.click(get_folder_path, outputs=resume) + gradient_checkpointing = gr.Checkbox( + label='Gradient checkpointing', value=False + ) + gradient_accumulation_steps = gr.Number( + label='Gradient accumulate steps', value='1' + ) with gr.Box(): with gr.Row(): create_caption = gr.Checkbox( @@ -770,6 +810,10 @@ def finetune_tab(): clip_skip, save_state, resume, + gradient_checkpointing, + gradient_accumulation_steps, + mem_eff_attn, + shuffle_caption, ] button_run.click(train_model, inputs=settings_list) diff --git a/lora_gui.py b/lora_gui.py index ea4e133..aa3e52d 100644 --- a/lora_gui.py +++ b/lora_gui.py @@ -74,6 +74,9 @@ def save_configuration( color_aug, flip_aug, clip_skip, + gradient_accumulation_steps, + mem_eff_attn, + # vae, ): original_file_path = file_path @@ -132,6 +135,9 @@ def save_configuration( 'color_aug': color_aug, 'flip_aug': flip_aug, 'clip_skip': clip_skip, + 'gradient_accumulation_steps': gradient_accumulation_steps, + 'mem_eff_attn': mem_eff_attn, + # 'vae': vae, } # Save the data to the selected file @@ -181,6 +187,9 @@ def open_configuration( color_aug, flip_aug, clip_skip, + gradient_accumulation_steps, + mem_eff_attn, + # vae, ): original_file_path = file_path @@ -240,6 +249,9 @@ def open_configuration( my_data.get('color_aug', color_aug), my_data.get('flip_aug', flip_aug), my_data.get('clip_skip', clip_skip), + my_data.get('gradient_accumulation_steps', gradient_accumulation_steps), + my_data.get('mem_eff_attn', mem_eff_attn), + # my_data.get('vae', vae), ) @@ -282,6 +294,9 @@ def train_model( color_aug, flip_aug, clip_skip, + gradient_accumulation_steps, + mem_eff_attn, + # vae, ): def save_inference_file(output_dir, v2, v_parameterization): # Copy inference model for v2 if required @@ -425,6 +440,8 @@ def train_model( run_cmd += ' --color_aug' if flip_aug: run_cmd += ' --flip_aug' + if mem_eff_attn: + run_cmd += ' --mem_eff_attn' run_cmd += ( f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"' ) @@ -475,6 +492,10 @@ def train_model( run_cmd += f' --network_weights="{lora_network_weights}"' if int(clip_skip) > 1: run_cmd += f' --clip_skip={str(clip_skip)}' + if int(gradient_accumulation_steps) > 1: + run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}' + # if not vae == '': + # run_cmd += f' --vae="{vae}"' print(run_cmd) # Run the command @@ -851,16 +872,19 @@ def lora_tab( label='No token padding', value=False ) - gradient_checkpointing_input = gr.Checkbox( + gradient_checkpointing = gr.Checkbox( label='Gradient checkpointing', value=False ) + gradient_accumulation_steps = gr.Number( + label='Gradient accumulate steps', value='1' + ) shuffle_caption = gr.Checkbox( label='Shuffle caption', value=False ) with gr.Row(): - save_state = gr.Checkbox( - label='Save training state', value=False + prior_loss_weight = gr.Number( + label='Prior loss weight', value=1.0 ) color_aug = gr.Checkbox( label='Color augmentation', value=False @@ -874,16 +898,25 @@ def lora_tab( clip_skip = gr.Slider( label='Clip skip', value='1', minimum=1, maximum=12, step=1 ) + mem_eff_attn = gr.Checkbox( + label='Memory efficient attention', value=False + ) with gr.Row(): + save_state = gr.Checkbox( + label='Save training state', value=False + ) resume = gr.Textbox( label='Resume from saved training state', placeholder='path to "last-state" state folder to resume from', ) resume_button = gr.Button('📂', elem_id='open_folder_small') resume_button.click(get_folder_path, outputs=resume) - prior_loss_weight = gr.Number( - label='Prior loss weight', value=1.0 - ) + # vae = gr.Textbox( + # label='VAE', + # placeholder='(Optiona) path to checkpoint of vae to replace for training', + # ) + # vae_button = gr.Button('📂', elem_id='open_folder_small') + # vae_button.click(get_any_file_path, outputs=vae) with gr.Tab('Tools'): gr.Markdown( 'This section provide Dreambooth tools to help setup your dataset...' @@ -921,7 +954,7 @@ def lora_tab( cache_latent_input, caption_extention_input, enable_bucket_input, - gradient_checkpointing_input, + gradient_checkpointing, full_fp16_input, no_token_padding_input, stop_text_encoder_training_input, @@ -939,6 +972,9 @@ def lora_tab( color_aug, flip_aug, clip_skip, + gradient_accumulation_steps, + mem_eff_attn, + # vae, ] button_open_config.click(