- 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

@@ -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

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(

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)

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(
+                prior_loss_weight = gr.Number(
-                    label='Save training state', value=False
+                    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(
+                # vae = gr.Textbox(
-                    label='Prior loss weight', value=1.0
+                #     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(