StableDiffusion
/
KohyaSS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

- 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
This commit is contained in:
bmaltais 2023-01-08 20:55:41 -05:00
parent f1d53ae3f9
commit a4262c0a66
4 changed files with 105 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

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

14
dreambooth_gui.py
View File

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

44
finetune_gui.py
View File

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

50
lora_gui.py
View File

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