KohyaSS/dreambooth_gui.py

945 lines
27 KiB
Python
Raw Normal View History

# v1: initial release
2022-12-14 19:40:24 +00:00
# v2: add open and save folder icons
2022-12-15 23:19:35 +00:00
# v3: Add new Utilities tab for Dreambooth folder preparation
2022-12-16 16:04:06 +00:00
# v3.1: Adding captionning of images to utilities
2022-12-13 14:20:25 +00:00
import gradio as gr
import json
import math
import os
import subprocess
import pathlib
2022-12-21 14:05:06 +00:00
import argparse
2022-12-18 18:11:10 +00:00
from library.common_gui import (
2022-12-17 01:26:26 +00:00
get_folder_path,
remove_doublequote,
get_file_path,
2023-01-01 19:14:58 +00:00
get_any_file_path,
get_saveasfile_path,
color_aug_changed,
save_inference_file,
gradio_advanced_training,
run_cmd_advanced_training,
2023-01-16 00:59:40 +00:00
run_cmd_training,
gradio_training,
gradio_config,
gradio_source_model,
2023-03-05 15:34:09 +00:00
# set_legacy_8bitadam,
2023-03-04 23:56:22 +00:00
update_my_data,
2023-03-20 12:47:00 +00:00
check_if_model_exist,
2023-03-01 18:14:47 +00:00
)
from library.tensorboard_gui import (
2023-02-27 00:49:22 +00:00
gradio_tensorboard,
start_tensorboard,
stop_tensorboard,
2022-12-17 01:26:26 +00:00
)
2022-12-23 01:18:51 +00:00
from library.dreambooth_folder_creation_gui import (
gradio_dreambooth_folder_creation_tab,
)
2022-12-22 16:51:34 +00:00
from library.utilities import utilities_tab
from library.sampler_gui import sample_gradio_config, run_cmd_sample
from easygui import msgbox
2022-12-13 14:20:25 +00:00
folder_symbol = '\U0001f4c2' # 📂
refresh_symbol = '\U0001f504' # 🔄
save_style_symbol = '\U0001f4be' # 💾
document_symbol = '\U0001F4C4' # 📄
2022-12-13 14:20:25 +00:00
2022-12-15 12:48:29 +00:00
def save_configuration(
save_as,
2022-12-13 14:20:25 +00:00
file_path,
pretrained_model_name_or_path,
v2,
2022-12-13 19:59:33 +00:00
v_parameterization,
2022-12-13 14:20:25 +00:00
logging_dir,
train_data_dir,
reg_data_dir,
output_dir,
max_resolution,
learning_rate,
lr_scheduler,
lr_warmup,
train_batch_size,
epoch,
save_every_n_epochs,
mixed_precision,
save_precision,
seed,
num_cpu_threads_per_process,
2023-01-16 00:59:40 +00:00
cache_latents,
caption_extension,
2022-12-13 18:49:14 +00:00
enable_bucket,
gradient_checkpointing,
2022-12-14 02:21:59 +00:00
full_fp16,
no_token_padding,
stop_text_encoder_training,
2023-03-05 15:34:09 +00:00
# use_8bit_adam,
2022-12-14 02:21:59 +00:00
xformers,
save_model_as,
shuffle_caption,
save_state,
resume,
prior_loss_weight,
color_aug,
flip_aug,
clip_skip,
vae,
output_name,
2023-01-10 14:38:32 +00:00
max_token_length,
max_train_epochs,
2023-01-16 00:59:40 +00:00
max_data_loader_n_workers,
mem_eff_attn,
gradient_accumulation_steps,
2023-02-06 01:07:00 +00:00
model_list,
keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale,
random_crop,
bucket_reso_steps,
2023-03-02 00:24:11 +00:00
caption_dropout_every_n_epochs,
caption_dropout_rate,
optimizer,
2023-03-02 00:24:11 +00:00
optimizer_args,
noise_offset,
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
2023-03-20 12:47:00 +00:00
sample_prompts,
additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma,
2022-12-13 14:20:25 +00:00
):
# Get list of function parameters and values
parameters = list(locals().items())
2022-12-14 19:40:24 +00:00
original_file_path = file_path
2022-12-17 01:26:26 +00:00
save_as_bool = True if save_as.get('label') == 'True' else False
2022-12-15 12:48:29 +00:00
if save_as_bool:
2022-12-17 01:26:26 +00:00
print('Save as...')
file_path = get_saveasfile_path(file_path)
2022-12-15 12:48:29 +00:00
else:
2022-12-17 01:26:26 +00:00
print('Save...')
if file_path == None or file_path == '':
file_path = get_saveasfile_path(file_path)
2022-12-14 19:40:24 +00:00
# print(file_path)
2022-12-19 14:47:35 +00:00
if file_path == None or file_path == '':
2022-12-15 12:48:29 +00:00
return original_file_path # In case a file_path was provided and the user decide to cancel the open action
2022-12-14 19:40:24 +00:00
2022-12-13 14:20:25 +00:00
# Return the values of the variables as a dictionary
variables = {
name: value
for name, value in parameters # locals().items()
if name
not in [
'file_path',
'save_as',
]
2022-12-13 14:20:25 +00:00
}
2023-02-20 12:56:24 +00:00
# Extract the destination directory from the file path
destination_directory = os.path.dirname(file_path)
# Create the destination directory if it doesn't exist
if not os.path.exists(destination_directory):
os.makedirs(destination_directory)
2022-12-13 14:20:25 +00:00
# Save the data to the selected file
2022-12-17 01:26:26 +00:00
with open(file_path, 'w') as file:
json.dump(variables, file, indent=2)
2022-12-13 14:20:25 +00:00
2022-12-14 19:40:24 +00:00
return file_path
2022-12-15 12:48:29 +00:00
def open_configuration(
ask_for_file,
2022-12-14 19:40:24 +00:00
file_path,
pretrained_model_name_or_path,
v2,
v_parameterization,
logging_dir,
train_data_dir,
reg_data_dir,
output_dir,
max_resolution,
learning_rate,
lr_scheduler,
lr_warmup,
train_batch_size,
epoch,
save_every_n_epochs,
mixed_precision,
save_precision,
seed,
num_cpu_threads_per_process,
2023-01-16 00:59:40 +00:00
cache_latents,
caption_extension,
2022-12-14 19:40:24 +00:00
enable_bucket,
gradient_checkpointing,
full_fp16,
no_token_padding,
stop_text_encoder_training,
2023-03-05 15:34:09 +00:00
# use_8bit_adam,
2022-12-14 19:40:24 +00:00
xformers,
save_model_as,
shuffle_caption,
save_state,
resume,
prior_loss_weight,
color_aug,
flip_aug,
clip_skip,
vae,
output_name,
2023-01-10 14:38:32 +00:00
max_token_length,
max_train_epochs,
2023-01-16 00:59:40 +00:00
max_data_loader_n_workers,
mem_eff_attn,
gradient_accumulation_steps,
2023-02-06 01:07:00 +00:00
model_list,
keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale,
random_crop,
bucket_reso_steps,
2023-03-02 00:24:11 +00:00
caption_dropout_every_n_epochs,
caption_dropout_rate,
optimizer,
2023-03-02 00:24:11 +00:00
optimizer_args,
noise_offset,
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
2023-03-20 12:47:00 +00:00
sample_prompts,
additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma,
2022-12-14 19:40:24 +00:00
):
# Get list of function parameters and values
parameters = list(locals().items())
2023-03-20 12:47:00 +00:00
ask_for_file = True if ask_for_file.get('label') == 'True' else False
2022-12-14 19:40:24 +00:00
original_file_path = file_path
2023-03-20 12:47:00 +00:00
if ask_for_file:
file_path = get_file_path(file_path)
2022-12-13 14:20:25 +00:00
if not file_path == '' and not file_path == None:
2022-12-14 19:40:24 +00:00
# load variables from JSON file
2022-12-17 01:26:26 +00:00
with open(file_path, 'r') as f:
2023-03-02 22:51:17 +00:00
my_data = json.load(f)
2023-01-16 00:59:40 +00:00
print('Loading config...')
2023-03-02 00:02:04 +00:00
# Update values to fix deprecated use_8bit_adam checkbox and set appropriate optimizer if it is set to True
2023-03-04 23:56:22 +00:00
my_data = update_my_data(my_data)
2022-12-14 19:40:24 +00:00
else:
file_path = original_file_path # In case a file_path was provided and the user decide to cancel the open action
2023-03-02 22:51:17 +00:00
my_data = {}
2022-12-13 14:20:25 +00:00
values = [file_path]
for key, value in parameters:
# Set the value in the dictionary to the corresponding value in `my_data`, or the default value if not found
if not key in ['ask_for_file', 'file_path']:
2023-03-02 22:51:17 +00:00
values.append(my_data.get(key, value))
return tuple(values)
2023-01-16 00:59:40 +00:00
2022-12-13 14:20:25 +00:00
def train_model(
pretrained_model_name_or_path,
v2,
2022-12-13 19:59:33 +00:00
v_parameterization,
2022-12-13 14:20:25 +00:00
logging_dir,
train_data_dir,
reg_data_dir,
output_dir,
max_resolution,
learning_rate,
lr_scheduler,
lr_warmup,
train_batch_size,
epoch,
save_every_n_epochs,
mixed_precision,
save_precision,
seed,
num_cpu_threads_per_process,
2023-01-16 00:59:40 +00:00
cache_latents,
2023-01-01 19:14:58 +00:00
caption_extension,
2022-12-13 18:49:14 +00:00
enable_bucket,
gradient_checkpointing,
2022-12-14 02:21:59 +00:00
full_fp16,
no_token_padding,
stop_text_encoder_training_pct,
2023-03-05 15:34:09 +00:00
# use_8bit_adam,
2022-12-14 02:21:59 +00:00
xformers,
save_model_as,
shuffle_caption,
save_state,
resume,
prior_loss_weight,
color_aug,
flip_aug,
clip_skip,
vae,
output_name,
2023-01-10 14:38:32 +00:00
max_token_length,
max_train_epochs,
2023-01-16 00:59:40 +00:00
max_data_loader_n_workers,
mem_eff_attn,
gradient_accumulation_steps,
2023-01-16 00:59:40 +00:00
model_list, # Keep this. Yes, it is unused here but required given the common list used
2023-01-27 12:33:44 +00:00
keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale,
random_crop,
bucket_reso_steps,
2023-03-02 00:24:11 +00:00
caption_dropout_every_n_epochs,
caption_dropout_rate,
optimizer,
2023-03-02 00:24:11 +00:00
optimizer_args,
noise_offset,
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
2023-03-20 12:47:00 +00:00
sample_prompts,
additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma,
2022-12-13 14:20:25 +00:00
):
2022-12-17 01:26:26 +00:00
if pretrained_model_name_or_path == '':
msgbox('Source model information is missing')
2022-12-16 00:04:26 +00:00
return
2022-12-17 01:26:26 +00:00
if train_data_dir == '':
msgbox('Image folder path is missing')
2022-12-16 00:04:26 +00:00
return
if not os.path.exists(train_data_dir):
2022-12-17 01:26:26 +00:00
msgbox('Image folder does not exist')
2022-12-16 00:04:26 +00:00
return
2022-12-17 01:26:26 +00:00
if reg_data_dir != '':
2022-12-16 00:04:26 +00:00
if not os.path.exists(reg_data_dir):
2022-12-17 01:26:26 +00:00
msgbox('Regularisation folder does not exist')
2022-12-16 00:04:26 +00:00
return
2022-12-17 01:26:26 +00:00
if output_dir == '':
msgbox('Output folder path is missing')
2022-12-16 00:04:26 +00:00
return
2023-03-20 12:47:00 +00:00
if check_if_model_exist(output_name, output_dir, save_model_as):
return
2023-03-22 00:34:21 +00:00
# Get a list of all subfolders in train_data_dir, excluding hidden folders
2022-12-14 02:21:59 +00:00
subfolders = [
2022-12-17 01:26:26 +00:00
f
for f in os.listdir(train_data_dir)
2023-03-28 15:54:42 +00:00
if os.path.isdir(os.path.join(train_data_dir, f))
and not f.startswith('.')
2022-12-14 02:21:59 +00:00
]
2022-12-13 14:20:25 +00:00
2023-03-22 16:55:30 +00:00
# Check if subfolders are present. If not let the user know and return
if not subfolders:
2023-03-28 15:54:42 +00:00
print(
'\033[33mNo subfolders were found in',
train_data_dir,
" can't train\...033[0m",
)
2023-03-22 16:55:30 +00:00
return
2022-12-13 14:20:25 +00:00
total_steps = 0
# Loop through each subfolder and extract the number of repeats
for folder in subfolders:
# Extract the number of repeats from the folder name
2023-03-22 16:55:30 +00:00
try:
repeats = int(folder.split('_')[0])
except ValueError:
2023-03-28 15:54:42 +00:00
print(
'\033[33mSubfolder',
folder,
"does not have a proper repeat value, please correct the name or remove it... can't train...\033[0m",
)
2023-03-22 16:55:30 +00:00
continue
2022-12-13 14:20:25 +00:00
# Count the number of images in the folder
2022-12-17 01:26:26 +00:00
num_images = len(
[
f
for f, lower_f in (
2023-03-28 15:54:42 +00:00
(file, file.lower())
for file in os.listdir(
os.path.join(train_data_dir, folder)
)
)
if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
2022-12-17 01:26:26 +00:00
]
)
2023-03-28 15:54:42 +00:00
2023-03-22 16:55:30 +00:00
if num_images == 0:
print(f'{folder} folder contain no images, skipping...')
else:
# Calculate the total number of steps for this folder
steps = repeats * num_images
total_steps += steps
# Print the result
print('\033[33mFolder', folder, ':', steps, 'steps\033[0m')
if total_steps == 0:
2023-03-28 15:54:42 +00:00
print(
'\033[33mNo images were found in folder',
train_data_dir,
'... please rectify!\033[0m',
)
2023-03-22 16:55:30 +00:00
return
2022-12-13 14:20:25 +00:00
# Print the result
# print(f"{total_steps} total steps")
2022-12-17 01:26:26 +00:00
if reg_data_dir == '':
reg_factor = 1
else:
2023-03-28 15:54:42 +00:00
print(
'\033[94mRegularisation images are used... Will double the number of steps required...\033[0m'
)
reg_factor = 2
2022-12-13 14:20:25 +00:00
# calculate max_train_steps
max_train_steps = int(
2022-12-14 02:21:59 +00:00
math.ceil(
2022-12-17 01:26:26 +00:00
float(total_steps)
/ int(train_batch_size)
* int(epoch)
* int(reg_factor)
)
)
print(f'max_train_steps = {max_train_steps}')
2022-12-13 14:20:25 +00:00
2022-12-14 02:21:59 +00:00
# calculate stop encoder training
if int(stop_text_encoder_training_pct) == -1:
stop_text_encoder_training = -1
elif stop_text_encoder_training_pct == None:
2022-12-14 02:21:59 +00:00
stop_text_encoder_training = 0
else:
stop_text_encoder_training = math.ceil(
2022-12-17 01:26:26 +00:00
float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
)
print(f'stop_text_encoder_training = {stop_text_encoder_training}')
2022-12-14 02:21:59 +00:00
2022-12-13 14:20:25 +00:00
lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
2022-12-17 01:26:26 +00:00
print(f'lr_warmup_steps = {lr_warmup_steps}')
2022-12-13 14:20:25 +00:00
run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_db.py"'
2022-12-13 14:20:25 +00:00
if v2:
2022-12-17 01:26:26 +00:00
run_cmd += ' --v2'
2022-12-13 19:59:33 +00:00
if v_parameterization:
2022-12-17 01:26:26 +00:00
run_cmd += ' --v_parameterization'
if enable_bucket:
2022-12-17 01:26:26 +00:00
run_cmd += ' --enable_bucket'
2022-12-14 02:21:59 +00:00
if no_token_padding:
2022-12-17 01:26:26 +00:00
run_cmd += ' --no_token_padding'
run_cmd += (
f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
2022-12-17 01:26:26 +00:00
)
2022-12-15 23:19:35 +00:00
run_cmd += f' --train_data_dir="{train_data_dir}"'
if len(reg_data_dir):
run_cmd += f' --reg_data_dir="{reg_data_dir}"'
2022-12-17 01:26:26 +00:00
run_cmd += f' --resolution={max_resolution}'
run_cmd += f' --output_dir="{output_dir}"'
run_cmd += f' --logging_dir="{logging_dir}"'
2022-12-19 15:39:04 +00:00
if not stop_text_encoder_training == 0:
run_cmd += (
f' --stop_text_encoder_training={stop_text_encoder_training}'
)
2022-12-18 01:36:31 +00:00
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}'
if not vae == '':
run_cmd += f' --vae="{vae}"'
if not output_name == '':
run_cmd += f' --output_name="{output_name}"'
2023-01-16 00:59:40 +00:00
if int(max_token_length) > 75:
2023-01-10 14:38:32 +00:00
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 == '':
2023-01-16 00:59:40 +00:00
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)}'
2023-01-16 00:59:40 +00:00
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,
2023-02-23 01:32:57 +00:00
optimizer=optimizer,
optimizer_args=optimizer_args,
2023-01-16 00:59:40 +00:00
)
run_cmd += run_cmd_advanced_training(
max_train_epochs=max_train_epochs,
max_data_loader_n_workers=max_data_loader_n_workers,
max_token_length=max_token_length,
resume=resume,
save_state=save_state,
mem_eff_attn=mem_eff_attn,
clip_skip=clip_skip,
flip_aug=flip_aug,
color_aug=color_aug,
shuffle_caption=shuffle_caption,
gradient_checkpointing=gradient_checkpointing,
full_fp16=full_fp16,
xformers=xformers,
2023-03-05 15:34:09 +00:00
# use_8bit_adam=use_8bit_adam,
2023-01-27 12:33:44 +00:00
keep_tokens=keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers=persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale=bucket_no_upscale,
random_crop=random_crop,
bucket_reso_steps=bucket_reso_steps,
2023-02-08 01:58:35 +00:00
caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
caption_dropout_rate=caption_dropout_rate,
2023-02-23 01:32:57 +00:00
noise_offset=noise_offset,
2023-03-07 12:42:13 +00:00
additional_parameters=additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size=vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma=min_snr_gamma,
)
2022-12-13 14:20:25 +00:00
run_cmd += run_cmd_sample(
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
sample_prompts,
2023-03-08 12:30:14 +00:00
output_dir,
)
2022-12-13 14:20:25 +00:00
print(run_cmd)
2022-12-13 14:20:25 +00:00
# Run the command
2023-03-05 16:43:59 +00:00
if os.name == 'posix':
os.system(run_cmd)
else:
subprocess.run(run_cmd)
2022-12-13 14:20:25 +00:00
2022-12-18 01:36:31 +00:00
# check if output_dir/last is a folder... therefore it is a diffuser model
last_dir = pathlib.Path(f'{output_dir}/{output_name}')
2022-12-18 01:36:31 +00:00
if not last_dir.is_dir():
2022-12-13 14:20:25 +00:00
# Copy inference model for v2 if required
save_inference_file(output_dir, v2, v_parameterization, output_name)
2022-12-13 14:20:25 +00:00
2022-12-22 16:51:34 +00:00
def dreambooth_tab(
2023-01-09 13:08:47 +00:00
train_data_dir=gr.Textbox(),
reg_data_dir=gr.Textbox(),
output_dir=gr.Textbox(),
logging_dir=gr.Textbox(),
2022-12-22 16:51:34 +00:00
):
dummy_db_true = gr.Label(value=True, visible=False)
dummy_db_false = gr.Label(value=False, visible=False)
2022-12-23 01:18:51 +00:00
gr.Markdown('Train a custom model using kohya dreambooth python code...')
2023-01-16 00:59:40 +00:00
(
button_open_config,
button_save_config,
button_save_as_config,
config_file_name,
button_load_config,
2023-01-16 00:59:40 +00:00
) = gradio_config()
(
pretrained_model_name_or_path,
v2,
v_parameterization,
save_model_as,
model_list,
) = gradio_source_model()
2022-12-23 01:18:51 +00:00
with gr.Tab('Folders'):
2022-12-22 16:51:34 +00:00
with gr.Row():
2023-01-09 13:08:47 +00:00
train_data_dir = gr.Textbox(
2022-12-22 16:51:34 +00:00
label='Image folder',
placeholder='Folder where the training folders containing the images are located',
)
train_data_dir_input_folder = gr.Button(
'📂', elem_id='open_folder_small'
)
train_data_dir_input_folder.click(
2023-03-04 23:56:22 +00:00
get_folder_path,
outputs=train_data_dir,
show_progress=False,
2022-12-22 16:51:34 +00:00
)
2023-01-09 13:08:47 +00:00
reg_data_dir = gr.Textbox(
2022-12-22 16:51:34 +00:00
label='Regularisation folder',
placeholder='(Optional) Folder where where the regularization folders containing the images are located',
)
reg_data_dir_input_folder = gr.Button(
'📂', elem_id='open_folder_small'
)
reg_data_dir_input_folder.click(
2023-03-04 23:56:22 +00:00
get_folder_path,
outputs=reg_data_dir,
show_progress=False,
2022-12-22 16:51:34 +00:00
)
with gr.Row():
2023-01-09 13:08:47 +00:00
output_dir = gr.Textbox(
label='Model output folder',
2022-12-22 16:51:34 +00:00
placeholder='Folder to output trained model',
)
output_dir_input_folder = gr.Button(
'📂', elem_id='open_folder_small'
)
output_dir_input_folder.click(get_folder_path, outputs=output_dir)
2023-01-09 13:08:47 +00:00
logging_dir = gr.Textbox(
2022-12-22 16:51:34 +00:00
label='Logging folder',
placeholder='Optional: enable logging and output TensorBoard log to this folder',
)
logging_dir_input_folder = gr.Button(
'📂', elem_id='open_folder_small'
)
logging_dir_input_folder.click(
2023-03-04 23:56:22 +00:00
get_folder_path,
outputs=logging_dir,
show_progress=False,
2022-12-22 16:51:34 +00:00
)
with gr.Row():
output_name = gr.Textbox(
label='Model output name',
placeholder='Name of the model to output',
value='last',
interactive=True,
)
2023-01-09 13:08:47 +00:00
train_data_dir.change(
2022-12-22 16:51:34 +00:00
remove_doublequote,
2023-01-09 13:08:47 +00:00
inputs=[train_data_dir],
outputs=[train_data_dir],
2022-12-21 14:05:06 +00:00
)
2023-01-09 13:08:47 +00:00
reg_data_dir.change(
2022-12-22 16:51:34 +00:00
remove_doublequote,
2023-01-09 13:08:47 +00:00
inputs=[reg_data_dir],
outputs=[reg_data_dir],
2022-12-21 14:05:06 +00:00
)
2023-01-09 13:08:47 +00:00
output_dir.change(
2022-12-22 16:51:34 +00:00
remove_doublequote,
2023-01-09 13:08:47 +00:00
inputs=[output_dir],
outputs=[output_dir],
2022-12-22 16:51:34 +00:00
)
2023-01-09 13:08:47 +00:00
logging_dir.change(
2022-12-22 16:51:34 +00:00
remove_doublequote,
2023-01-09 13:08:47 +00:00
inputs=[logging_dir],
outputs=[logging_dir],
2022-12-21 14:05:06 +00:00
)
2022-12-22 16:51:34 +00:00
with gr.Tab('Training parameters'):
2023-01-16 00:59:40 +00:00
(
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,
2023-03-02 00:24:11 +00:00
optimizer,
optimizer_args,
2023-01-16 00:59:40 +00:00
) = gradio_training(
learning_rate_value='1e-5',
lr_scheduler_value='cosine',
lr_warmup_value='10',
)
2022-12-22 16:51:34 +00:00
with gr.Row():
2023-01-09 13:08:47 +00:00
max_resolution = gr.Textbox(
2022-12-22 16:51:34 +00:00
label='Max resolution',
value='512,512',
placeholder='512,512',
)
2023-01-09 13:08:47 +00:00
stop_text_encoder_training = gr.Slider(
minimum=-1,
2022-12-22 16:51:34 +00:00
maximum=100,
value=0,
step=1,
label='Stop text encoder training',
)
enable_bucket = gr.Checkbox(label='Enable buckets', value=True)
2022-12-22 16:51:34 +00:00
with gr.Accordion('Advanced Configuration', open=False):
with gr.Row():
2023-01-09 13:08:47 +00:00
no_token_padding = gr.Checkbox(
2022-12-22 16:51:34 +00:00
label='No token padding', value=False
)
gradient_accumulation_steps = gr.Number(
label='Gradient accumulate steps', value='1'
)
2022-12-22 16:51:34 +00:00
with gr.Row():
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')
2023-03-04 23:56:22 +00:00
vae_button.click(
get_any_file_path,
outputs=vae,
show_progress=False,
)
(
2023-03-05 15:34:09 +00:00
# use_8bit_adam,
xformers,
full_fp16,
gradient_checkpointing,
shuffle_caption,
color_aug,
flip_aug,
clip_skip,
mem_eff_attn,
save_state,
resume,
max_token_length,
max_train_epochs,
max_data_loader_n_workers,
2023-01-27 12:33:44 +00:00
keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale,
random_crop,
bucket_reso_steps,
2023-03-02 00:24:11 +00:00
caption_dropout_every_n_epochs,
caption_dropout_rate,
noise_offset,
2023-03-07 12:42:13 +00:00
additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma,
) = gradio_advanced_training()
color_aug.change(
color_aug_changed,
inputs=[color_aug],
2023-01-16 00:59:40 +00:00
outputs=[cache_latents],
)
(
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
sample_prompts,
) = sample_gradio_config()
2022-12-23 01:18:51 +00:00
with gr.Tab('Tools'):
gr.Markdown(
'This section provide Dreambooth tools to help setup your dataset...'
)
2022-12-23 01:18:51 +00:00
gradio_dreambooth_folder_creation_tab(
2023-01-09 13:08:47 +00:00
train_data_dir_input=train_data_dir,
reg_data_dir_input=reg_data_dir,
output_dir_input=output_dir,
logging_dir_input=logging_dir,
2022-12-23 01:18:51 +00:00
)
2022-12-22 16:51:34 +00:00
2023-02-27 00:49:22 +00:00
button_run = gr.Button('Train model', variant='primary')
2023-03-02 00:24:11 +00:00
2023-02-27 00:49:22 +00:00
# Setup gradio tensorboard buttons
button_start_tensorboard, button_stop_tensorboard = gradio_tensorboard()
2023-03-02 00:24:11 +00:00
2023-02-27 00:49:22 +00:00
button_start_tensorboard.click(
start_tensorboard,
inputs=logging_dir,
2023-03-04 23:56:22 +00:00
show_progress=False,
2023-02-27 00:49:22 +00:00
)
2023-03-02 00:24:11 +00:00
2023-02-27 00:49:22 +00:00
button_stop_tensorboard.click(
stop_tensorboard,
2023-03-04 23:56:22 +00:00
show_progress=False,
2023-02-27 00:49:22 +00:00
)
settings_list = [
2023-01-09 13:08:47 +00:00
pretrained_model_name_or_path,
v2,
v_parameterization,
logging_dir,
train_data_dir,
reg_data_dir,
output_dir,
max_resolution,
learning_rate,
lr_scheduler,
lr_warmup,
train_batch_size,
epoch,
save_every_n_epochs,
mixed_precision,
save_precision,
seed,
num_cpu_threads_per_process,
2023-01-16 00:59:40 +00:00
cache_latents,
caption_extension,
2023-01-09 13:08:47 +00:00
enable_bucket,
gradient_checkpointing,
full_fp16,
no_token_padding,
stop_text_encoder_training,
2023-03-05 15:34:09 +00:00
# use_8bit_adam,
2023-01-09 13:08:47 +00:00
xformers,
save_model_as,
shuffle_caption,
save_state,
resume,
prior_loss_weight,
color_aug,
flip_aug,
clip_skip,
vae,
output_name,
2023-01-10 14:38:32 +00:00
max_token_length,
max_train_epochs,
2023-01-16 00:59:40 +00:00
max_data_loader_n_workers,
mem_eff_attn,
gradient_accumulation_steps,
2023-01-16 00:59:40 +00:00
model_list,
2023-01-27 12:33:44 +00:00
keep_tokens,
2023-02-04 16:55:06 +00:00
persistent_data_loader_workers,
2023-02-06 01:07:00 +00:00
bucket_no_upscale,
random_crop,
bucket_reso_steps,
2023-03-02 00:24:11 +00:00
caption_dropout_every_n_epochs,
caption_dropout_rate,
optimizer,
optimizer_args,
noise_offset,
sample_every_n_steps,
sample_every_n_epochs,
sample_sampler,
sample_prompts,
2023-03-07 12:42:13 +00:00
additional_parameters,
2023-03-22 00:20:57 +00:00
vae_batch_size,
2023-03-28 15:54:42 +00:00
min_snr_gamma,
]
2022-12-22 16:51:34 +00:00
button_open_config.click(
open_configuration,
inputs=[dummy_db_true, config_file_name] + settings_list,
outputs=[config_file_name] + settings_list,
show_progress=False,
)
2023-03-20 12:47:00 +00:00
button_load_config.click(
open_configuration,
inputs=[dummy_db_false, config_file_name] + settings_list,
outputs=[config_file_name] + settings_list,
2023-03-04 23:56:22 +00:00
show_progress=False,
2022-12-22 16:51:34 +00:00
)
button_save_config.click(
save_configuration,
inputs=[dummy_db_false, config_file_name] + settings_list,
2022-12-22 16:51:34 +00:00
outputs=[config_file_name],
2023-03-04 23:56:22 +00:00
show_progress=False,
2022-12-22 16:51:34 +00:00
)
button_save_as_config.click(
save_configuration,
inputs=[dummy_db_true, config_file_name] + settings_list,
2022-12-22 16:51:34 +00:00
outputs=[config_file_name],
2023-03-04 23:56:22 +00:00
show_progress=False,
2022-12-22 16:51:34 +00:00
)
button_run.click(
train_model,
inputs=settings_list,
2023-03-04 23:56:22 +00:00
show_progress=False,
2022-12-22 16:51:34 +00:00
)
return (
2023-01-09 13:08:47 +00:00
train_data_dir,
reg_data_dir,
output_dir,
logging_dir,
2022-12-22 16:51:34 +00:00
)
2022-12-21 14:05:06 +00:00
def UI(**kwargs):
css = ''
if os.path.exists('./style.css'):
with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
print('Load CSS...')
css += file.read() + '\n'
interface = gr.Blocks(css=css)
with interface:
with gr.Tab('Dreambooth'):
(
train_data_dir_input,
reg_data_dir_input,
output_dir_input,
logging_dir_input,
) = dreambooth_tab()
with gr.Tab('Utilities'):
utilities_tab(
train_data_dir_input=train_data_dir_input,
reg_data_dir_input=reg_data_dir_input,
output_dir_input=output_dir_input,
logging_dir_input=logging_dir_input,
enable_copy_info_button=True,
)
# Show the interface
2023-03-02 00:24:11 +00:00
launch_kwargs = {}
if not kwargs.get('username', None) == '':
2023-03-02 00:24:11 +00:00
launch_kwargs['auth'] = (
kwargs.get('username', None),
kwargs.get('password', None),
)
if kwargs.get('server_port', 0) > 0:
2023-03-02 00:24:11 +00:00
launch_kwargs['server_port'] = kwargs.get('server_port', 0)
if kwargs.get('inbrowser', False):
launch_kwargs['inbrowser'] = kwargs.get('inbrowser', False)
print(launch_kwargs)
interface.launch(**launch_kwargs)
2023-03-02 00:24:11 +00:00
2022-12-21 14:05:06 +00:00
if __name__ == '__main__':
2022-12-22 16:51:34 +00:00
# torch.cuda.set_per_process_memory_fraction(0.48)
parser = argparse.ArgumentParser()
parser.add_argument(
'--username', type=str, default='', help='Username for authentication'
)
parser.add_argument(
'--password', type=str, default='', help='Password for authentication'
)
parser.add_argument(
2023-03-02 00:24:11 +00:00
'--server_port',
type=int,
default=0,
help='Port to run the server listener on',
)
parser.add_argument(
'--inbrowser', action='store_true', help='Open in browser'
)
2022-12-22 16:51:34 +00:00
args = parser.parse_args()
2023-03-02 00:24:11 +00:00
UI(
username=args.username,
password=args.password,
inbrowser=args.inbrowser,
server_port=args.server_port,
)