Complete training code refactoring
This commit is contained in:
bmaltais
parent
123cf4e3c5
commit
567221549f
@ -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(
|
||||
|
||||
369
finetune_gui.py
369
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)
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
238
lora_gui.py
238
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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user