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

Adding support for new parameters

Browse Source
This commit is contained in:
bmaltais 2022-12-13 21:21:59 -05:00
parent 71765ae243
commit 469b15b579
6 changed files with 184 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.vscode/settings.json vendored Normal file
View File

@ -0,0 +1,3 @@
{
"python.linting.enabled": true
}

BIN
None/20221213174728/dreambooth/events.out.tfevents.1670971669.DESKTOP-4M9DSE4.26752.0 Normal file
View File

Binary file not shown.

BIN
None/20221213174915/dreambooth/events.out.tfevents.1670971772.DESKTOP-4M9DSE4.19112.0 Normal file
View File

Binary file not shown.

BIN
None/20221213182833/dreambooth/events.out.tfevents.1670974130.DESKTOP-4M9DSE4.3624.0 Normal file
View File

Binary file not shown.

BIN
None/20221213203916/dreambooth/events.out.tfevents.1670981976.DESKTOP-4M9DSE4.27240.0 Normal file
View File

Binary file not shown.

235
dreambooth_gui.py
View File

@ -38,7 +38,11 @@ def save_variables(
use_safetensors, use_safetensors,
enable_bucket, enable_bucket,
gradient_checkpointing, gradient_checkpointing,
full_fp16 full_fp16,
no_token_padding,
stop_text_encoder_training,
use_8bit_adam,
xformers,
): ):
# Return the values of the variables as a dictionary # Return the values of the variables as a dictionary
variables = { variables = {
@ -67,7 +71,11 @@ def save_variables(
"use_safetensors": use_safetensors, "use_safetensors": use_safetensors,
"enable_bucket": enable_bucket, "enable_bucket": enable_bucket,
"gradient_checkpointing": gradient_checkpointing, "gradient_checkpointing": gradient_checkpointing,
"full_fp16": full_fp16 "full_fp16": full_fp16,
"no_token_padding": no_token_padding,
"stop_text_encoder_training": stop_text_encoder_training,
"use_8bit_adam": use_8bit_adam,
"xformers": xformers,
} }
# Save the data to the selected file # Save the data to the selected file
@ -108,6 +116,10 @@ def load_variables(file_path):
my_data.get("enable_bucket", None), my_data.get("enable_bucket", None),
my_data.get("gradient_checkpointing", None), my_data.get("gradient_checkpointing", None),
my_data.get("full_fp16", None), my_data.get("full_fp16", None),
my_data.get("no_token_padding", None),
my_data.get("stop_text_encoder_training", None),
my_data.get("use_8bit_adam", None),
my_data.get("xformers", None),
) )
@ -137,7 +149,11 @@ def train_model(
use_safetensors, use_safetensors,
enable_bucket, enable_bucket,
gradient_checkpointing, gradient_checkpointing,
full_fp16 full_fp16,
no_token_padding,
stop_text_encoder_training_pct,
use_8bit_adam,
xformers,
): ):
def save_inference_file(output_dir, v2, v_parameterization): def save_inference_file(output_dir, v2, v_parameterization):
# Copy inference model for v2 if required # Copy inference model for v2 if required
@ -155,8 +171,11 @@ def train_model(
) )
# Get a list of all subfolders in train_data_dir # Get a list of all subfolders in train_data_dir
subfolders = [f for f in os.listdir(train_data_dir) if os.path.isdir( subfolders = [
os.path.join(train_data_dir, f))] f
for f in os.listdir(train_data_dir)
if os.path.isdir(os.path.join(train_data_dir, f))
]
total_steps = 0 total_steps = 0
@ -166,8 +185,16 @@ def train_model(
repeats = int(folder.split("_")[0]) repeats = int(folder.split("_")[0])
# Count the number of images in the folder # Count the number of images in the folder
num_images = len([f for f in os.listdir(os.path.join(train_data_dir, folder)) if f.endswith( num_images = len(
".jpg") or f.endswith(".jpeg") or f.endswith(".png") or f.endswith(".webp")]) [
f
for f in os.listdir(os.path.join(train_data_dir, folder))
if f.endswith(".jpg")
or f.endswith(".jpeg")
or f.endswith(".png")
or f.endswith(".webp")
]
)
# Calculate the total number of steps for this folder # Calculate the total number of steps for this folder
steps = repeats * num_images steps = repeats * num_images
@ -182,15 +209,28 @@ def train_model(
if reg_data_dir == "": if reg_data_dir == "":
reg_factor = 1 reg_factor = 1
else: else:
print("Regularisation images are used... Will double the number of steps required...") print(
"Regularisation images are used... Will double the number of steps required..."
)
reg_factor = 2 reg_factor = 2
# calculate max_train_steps # calculate max_train_steps
max_train_steps = int( max_train_steps = int(
math.ceil(float(total_steps) / int(train_batch_size) * int(epoch) * int(reg_factor)) math.ceil(
float(total_steps) / int(train_batch_size) * int(epoch) * int(reg_factor)
)
) )
print(f"max_train_steps = {max_train_steps}") print(f"max_train_steps = {max_train_steps}")
# calculate stop encoder training
if stop_text_encoder_training_pct == None:
stop_text_encoder_training = 0
else:
stop_text_encoder_training = math.ceil(
float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
)
print(f"stop_text_encoder_training = {stop_text_encoder_training}")
lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100)) lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
print(f"lr_warmup_steps = {lr_warmup_steps}") print(f"lr_warmup_steps = {lr_warmup_steps}")
@ -209,6 +249,12 @@ def train_model(
run_cmd += " --gradient_checkpointing" run_cmd += " --gradient_checkpointing"
if full_fp16: if full_fp16:
run_cmd += " --full_fp16" run_cmd += " --full_fp16"
if no_token_padding:
run_cmd += " --no_token_padding"
if use_8bit_adam:
run_cmd += " --use_8bit_adam"
if xformers:
run_cmd += " --xformers"
run_cmd += f" --pretrained_model_name_or_path={pretrained_model_name_or_path}" run_cmd += f" --pretrained_model_name_or_path={pretrained_model_name_or_path}"
run_cmd += f" --train_data_dir={train_data_dir}" run_cmd += f" --train_data_dir={train_data_dir}"
run_cmd += f" --reg_data_dir={reg_data_dir}" run_cmd += f" --reg_data_dir={reg_data_dir}"
@ -227,6 +273,7 @@ def train_model(
run_cmd += f" --save_precision={save_precision}" run_cmd += f" --save_precision={save_precision}"
run_cmd += f" --logging_dir={logging_dir}" run_cmd += f" --logging_dir={logging_dir}"
run_cmd += f" --caption_extention={caption_extention}" run_cmd += f" --caption_extention={caption_extention}"
run_cmd += f" --stop_text_encoder_training={stop_text_encoder_training}"
print(run_cmd) print(run_cmd)
# Run the command # Run the command
@ -245,8 +292,7 @@ def train_model(
save_inference_file(output_dir, v2, v_parameterization) save_inference_file(output_dir, v2, v_parameterization)
if convert_to_safetensors: if convert_to_safetensors:
print( print(f"Converting diffuser model {last_dir} to {last_dir}.safetensors")
f"Converting diffuser model {last_dir} to {last_dir}.safetensors")
os.system( os.system(
f"python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.safetensors --{save_precision}" f"python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.safetensors --{save_precision}"
) )
@ -262,7 +308,10 @@ def train_model(
def set_pretrained_model_name_or_path_input(value, v2, v_parameterization): def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
# define a list of substrings to search for # define a list of substrings to search for
substrings_v2 = ["stabilityai/stable-diffusion-2-1-base", "stabilityai/stable-diffusion-2-base"] substrings_v2 = [
"stabilityai/stable-diffusion-2-1-base",
"stabilityai/stable-diffusion-2-base",
]
# check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v2 list # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v2 list
if str(value) in substrings_v2: if str(value) in substrings_v2:
@ -273,18 +322,26 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
return value, v2, v_parameterization return value, v2, v_parameterization
# define a list of substrings to search for v-objective # define a list of substrings to search for v-objective
substrings_v_parameterization = ["stabilityai/stable-diffusion-2-1", "stabilityai/stable-diffusion-2"] substrings_v_parameterization = [
"stabilityai/stable-diffusion-2-1",
"stabilityai/stable-diffusion-2",
]
# check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v_parameterization list # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v_parameterization list
if str(value) in substrings_v_parameterization: if str(value) in substrings_v_parameterization:
print("SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization") print(
"SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization"
)
v2 = True v2 = True
v_parameterization = True v_parameterization = True
return value, v2, v_parameterization return value, v2, v_parameterization
# define a list of substrings to v1.x # define a list of substrings to v1.x
substrings_v1_model = ["CompVis/stable-diffusion-v1-4", "runwayml/stable-diffusion-v1-5"] substrings_v1_model = [
"CompVis/stable-diffusion-v1-4",
"runwayml/stable-diffusion-v1-5",
]
if str(value) in substrings_v1_model: if str(value) in substrings_v1_model:
v2 = False v2 = False
@ -299,16 +356,29 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
return value, v2, v_parameterization return value, v2, v_parameterization
def remove_doublequote(file_path):
if file_path != None:
file_path = file_path.replace('"', '')
return file_path
interface = gr.Blocks() interface = gr.Blocks()
with interface: with interface:
gr.Markdown("Enter kohya finetuner parameter using this interface.") gr.Markdown("Enter kohya finetuner parameter using this interface.")
with gr.Accordion("Configuration File Load/Save", open=False): with gr.Accordion("Configuration File Load/Save", open=False):
with gr.Row(): with gr.Row():
config_file_name = gr.Textbox( config_file_name = gr.Textbox(label="Config file name")
label="Config file name")
button_load_config = gr.Button("Load config") button_load_config = gr.Button("Load config")
button_save_config = gr.Button("Save config") button_save_config = gr.Button("Save config")
config_file_name.change(
remove_doublequote,
inputs=[config_file_name],
outputs=[
config_file_name
]
)
with gr.Tab("Source model"): with gr.Tab("Source model"):
# Define the input elements # Define the input elements
with gr.Row(): with gr.Row():
@ -325,39 +395,81 @@ with interface:
"stabilityai/stable-diffusion-2-1", "stabilityai/stable-diffusion-2-1",
"stabilityai/stable-diffusion-2", "stabilityai/stable-diffusion-2",
"runwayml/stable-diffusion-v1-5", "runwayml/stable-diffusion-v1-5",
"CompVis/stable-diffusion-v1-4" "CompVis/stable-diffusion-v1-4",
], ],
) )
with gr.Row(): with gr.Row():
v2_input = gr.Checkbox(label="v2", value=True) v2_input = gr.Checkbox(label="v2", value=True)
v_parameterization_input = gr.Checkbox(label="v_parameterization", value=False) v_parameterization_input = gr.Checkbox(
label="v_parameterization", value=False
)
pretrained_model_name_or_path_input.change(
remove_doublequote,
inputs=[pretrained_model_name_or_path_input],
outputs=[
pretrained_model_name_or_path_input
]
)
model_list.change( model_list.change(
set_pretrained_model_name_or_path_input, set_pretrained_model_name_or_path_input,
inputs=[model_list, v2_input, v_parameterization_input], inputs=[model_list, v2_input, v_parameterization_input],
outputs=[pretrained_model_name_or_path_input, outputs=[
v2_input, v_parameterization_input], pretrained_model_name_or_path_input,
v2_input,
v_parameterization_input,
],
) )
with gr.Tab("Directories"): with gr.Tab("Directories"):
with gr.Row(): with gr.Row():
train_data_dir_input = gr.Textbox( train_data_dir_input = gr.Textbox(
label="Image folder", placeholder="Directory where the training folders containing the images are located" label="Image folder",
placeholder="Directory where the training folders containing the images are located",
) )
reg_data_dir_input = gr.Textbox( reg_data_dir_input = gr.Textbox(
label="Regularisation folder", placeholder="(Optional) Directory where where the regularization folders containing the images are located" label="Regularisation folder",
placeholder="(Optional) Directory where where the regularization folders containing the images are located",
) )
with gr.Row(): with gr.Row():
output_dir_input = gr.Textbox( output_dir_input = gr.Textbox(
label="Output directory", label="Output directory",
placeholder="Directory to output trained model", placeholder="Directory to output trained model",
) )
logging_dir_input = gr.Textbox( logging_dir_input = gr.Textbox(
label="Logging directory", placeholder="Optional: enable logging and output TensorBoard log to this directory" label="Logging directory",
placeholder="Optional: enable logging and output TensorBoard log to this directory",
) )
train_data_dir_input.change(
remove_doublequote,
inputs=[train_data_dir_input],
outputs=[
train_data_dir_input
]
)
reg_data_dir_input.change(
remove_doublequote,
inputs=[reg_data_dir_input],
outputs=[
reg_data_dir_input
]
)
output_dir_input.change(
remove_doublequote,
inputs=[output_dir_input],
outputs=[
output_dir_input
]
)
logging_dir_input.change(
remove_doublequote,
inputs=[logging_dir_input],
outputs=[
logging_dir_input
]
)
with gr.Tab("Training parameters"): with gr.Tab("Training parameters"):
with gr.Row(): with gr.Row():
learning_rate_input = gr.Textbox( learning_rate_input = gr.Textbox(label="Learning rate", value=1e-6)
label="Learning rate", value=1e-6)
lr_scheduler_input = gr.Dropdown( lr_scheduler_input = gr.Dropdown(
label="LR Scheduler", label="LR Scheduler",
choices=[ choices=[
@ -372,13 +484,11 @@ with interface:
) )
lr_warmup_input = gr.Textbox(label="LR warmup", value=0) lr_warmup_input = gr.Textbox(label="LR warmup", value=0)
with gr.Row(): with gr.Row():
train_batch_size_input = gr.Textbox( train_batch_size_input = gr.Slider(
label="Train batch size", value=1 minimum=1, maximum=32, label="Train batch size", value=1, step=1
) )
epoch_input = gr.Textbox(label="Epoch", value=1) epoch_input = gr.Textbox(label="Epoch", value=1)
save_every_n_epochs_input = gr.Textbox( save_every_n_epochs_input = gr.Textbox(label="Save every N epochs", value=1)
label="Save every N epochs", value=1
)
with gr.Row(): with gr.Row():
mixed_precision_input = gr.Dropdown( mixed_precision_input = gr.Dropdown(
label="Mixed precision", label="Mixed precision",
@ -398,41 +508,46 @@ with interface:
], ],
value="fp16", value="fp16",
) )
num_cpu_threads_per_process_input = gr.Textbox( num_cpu_threads_per_process_input = gr.Slider(
label="Number of CPU threads per process", value=4 minimum=1,
maximum=os.cpu_count(),
step=1,
label="Number of CPU threads per process",
value=os.cpu_count(),
) )
with gr.Row(): with gr.Row():
seed_input = gr.Textbox(label="Seed", value=1234) seed_input = gr.Textbox(label="Seed", value=1234)
max_resolution_input = gr.Textbox( max_resolution_input = gr.Textbox(label="Max resolution", value="512,512")
label="Max resolution", value="512,512" with gr.Row():
)
caption_extention_input = gr.Textbox( caption_extention_input = gr.Textbox(
label="Caption Extension", placeholder="(Optional) Extension for caption files. default: .caption") label="Caption Extension",
placeholder="(Optional) Extension for caption files. default: .caption",
)
stop_text_encoder_training_input = gr.Slider(minimum=0, maximum=100, value=0, step=1,
label="Stop text encoder training",
)
with gr.Row(): with gr.Row():
use_safetensors_input = gr.Checkbox( use_safetensors_input = gr.Checkbox(
label="Use safetensor when saving", value=False label="Use safetensor when saving", value=False
) )
enable_bucket_input = gr.Checkbox( enable_bucket_input = gr.Checkbox(label="Enable buckets", value=False)
label="Enable buckets", value=False cache_latent_input = gr.Checkbox(label="Cache latent", value=True)
)
cache_latent_input = gr.Checkbox(
label="Cache latent", value=True
)
gradient_checkpointing_input = gr.Checkbox( gradient_checkpointing_input = gr.Checkbox(
label="Gradient checkpointing", value=False label="Gradient checkpointing", value=False
) )
with gr.Row():
full_fp16_input = gr.Checkbox( full_fp16_input = gr.Checkbox(
label="Full fp16 training (experimental)", value=False label="Full fp16 training (experimental)", value=False
) )
no_token_padding_input = gr.Checkbox(label="No tokan padding", value=False)
use_8bit_adam_input = gr.Checkbox(label="Use 8bit adam", value=True)
xformers_input = gr.Checkbox(label="USe xformers", value=True)
with gr.Tab("Model conversion"): with gr.Tab("Model conversion"):
convert_to_safetensors_input = gr.Checkbox( convert_to_safetensors_input = gr.Checkbox(
label="Convert to SafeTensors", value=False label="Convert to SafeTensors", value=False
) )
convert_to_ckpt_input = gr.Checkbox( convert_to_ckpt_input = gr.Checkbox(label="Convert to CKPT", value=False)
label="Convert to CKPT", value=False
)
button_run = gr.Button("Run") button_run = gr.Button("Run")
@ -465,8 +580,12 @@ with interface:
use_safetensors_input, use_safetensors_input,
enable_bucket_input, enable_bucket_input,
gradient_checkpointing_input, gradient_checkpointing_input,
full_fp16_input full_fp16_input,
] no_token_padding_input,
stop_text_encoder_training_input,
use_8bit_adam_input,
xformers_input,
],
) )
button_save_config.click( button_save_config.click(
@ -498,8 +617,12 @@ with interface:
use_safetensors_input, use_safetensors_input,
enable_bucket_input, enable_bucket_input,
gradient_checkpointing_input, gradient_checkpointing_input,
full_fp16_input full_fp16_input,
] no_token_padding_input,
stop_text_encoder_training_input,
use_8bit_adam_input,
xformers_input,
],
) )
button_run.click( button_run.click(
train_model, train_model,
@ -529,8 +652,12 @@ with interface:
use_safetensors_input, use_safetensors_input,
enable_bucket_input, enable_bucket_input,
gradient_checkpointing_input, gradient_checkpointing_input,
full_fp16_input full_fp16_input,
] no_token_padding_input,
stop_text_encoder_training_input,
use_8bit_adam_input,
xformers_input,
],
) )
# Show the interface # Show the interface