Reformat code with blue
This commit is contained in:
bmaltais
parent
90dad5471c
commit
42a0732e0d
@ -10,12 +10,17 @@ import os
|
||||
import subprocess
|
||||
import pathlib
|
||||
import shutil
|
||||
from dreambooth_gui.dreambooth_folder_creation import gradio_dreambooth_folder_creation_tab
|
||||
from dreambooth_gui.dreambooth_folder_creation import (
|
||||
gradio_dreambooth_folder_creation_tab,
|
||||
)
|
||||
from dreambooth_gui.caption_gui import gradio_caption_gui_tab
|
||||
from dreambooth_gui.common_gui import get_folder_path, remove_doublequote, get_file_path
|
||||
from dreambooth_gui.common_gui import (
|
||||
get_folder_path,
|
||||
remove_doublequote,
|
||||
get_file_path,
|
||||
)
|
||||
from easygui import filesavebox, msgbox
|
||||
|
||||
# sys.path.insert(0, './dreambooth_gui')
|
||||
|
||||
def save_configuration(
|
||||
save_as,
|
||||
@ -53,22 +58,22 @@ def save_configuration(
|
||||
):
|
||||
original_file_path = file_path
|
||||
|
||||
save_as_bool = True if save_as.get("label") == "True" else False
|
||||
save_as_bool = True if save_as.get('label') == 'True' else False
|
||||
|
||||
if save_as_bool:
|
||||
print("Save as...")
|
||||
print('Save as...')
|
||||
file_path = filesavebox(
|
||||
"Select the config file to save",
|
||||
default="finetune.json",
|
||||
filetypes="*.json",
|
||||
'Select the config file to save',
|
||||
default='finetune.json',
|
||||
filetypes='*.json',
|
||||
)
|
||||
else:
|
||||
print("Save...")
|
||||
if file_path == None or file_path == "":
|
||||
print('Save...')
|
||||
if file_path == None or file_path == '':
|
||||
file_path = filesavebox(
|
||||
"Select the config file to save",
|
||||
default="finetune.json",
|
||||
filetypes="*.json",
|
||||
'Select the config file to save',
|
||||
default='finetune.json',
|
||||
filetypes='*.json',
|
||||
)
|
||||
|
||||
if file_path == None:
|
||||
@ -76,40 +81,40 @@ def save_configuration(
|
||||
|
||||
# Return the values of the variables as a dictionary
|
||||
variables = {
|
||||
"pretrained_model_name_or_path": pretrained_model_name_or_path,
|
||||
"v2": v2,
|
||||
"v_parameterization": v_parameterization,
|
||||
"logging_dir": logging_dir,
|
||||
"train_data_dir": train_data_dir,
|
||||
"reg_data_dir": reg_data_dir,
|
||||
"output_dir": output_dir,
|
||||
"max_resolution": max_resolution,
|
||||
"learning_rate": learning_rate,
|
||||
"lr_scheduler": lr_scheduler,
|
||||
"lr_warmup": lr_warmup,
|
||||
"train_batch_size": train_batch_size,
|
||||
"epoch": epoch,
|
||||
"save_every_n_epochs": save_every_n_epochs,
|
||||
"mixed_precision": mixed_precision,
|
||||
"save_precision": save_precision,
|
||||
"seed": seed,
|
||||
"num_cpu_threads_per_process": num_cpu_threads_per_process,
|
||||
"convert_to_safetensors": convert_to_safetensors,
|
||||
"convert_to_ckpt": convert_to_ckpt,
|
||||
"cache_latent": cache_latent,
|
||||
"caption_extention": caption_extention,
|
||||
"use_safetensors": use_safetensors,
|
||||
"enable_bucket": enable_bucket,
|
||||
"gradient_checkpointing": gradient_checkpointing,
|
||||
"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,
|
||||
'pretrained_model_name_or_path': pretrained_model_name_or_path,
|
||||
'v2': v2,
|
||||
'v_parameterization': v_parameterization,
|
||||
'logging_dir': logging_dir,
|
||||
'train_data_dir': train_data_dir,
|
||||
'reg_data_dir': reg_data_dir,
|
||||
'output_dir': output_dir,
|
||||
'max_resolution': max_resolution,
|
||||
'learning_rate': learning_rate,
|
||||
'lr_scheduler': lr_scheduler,
|
||||
'lr_warmup': lr_warmup,
|
||||
'train_batch_size': train_batch_size,
|
||||
'epoch': epoch,
|
||||
'save_every_n_epochs': save_every_n_epochs,
|
||||
'mixed_precision': mixed_precision,
|
||||
'save_precision': save_precision,
|
||||
'seed': seed,
|
||||
'num_cpu_threads_per_process': num_cpu_threads_per_process,
|
||||
'convert_to_safetensors': convert_to_safetensors,
|
||||
'convert_to_ckpt': convert_to_ckpt,
|
||||
'cache_latent': cache_latent,
|
||||
'caption_extention': caption_extention,
|
||||
'use_safetensors': use_safetensors,
|
||||
'enable_bucket': enable_bucket,
|
||||
'gradient_checkpointing': gradient_checkpointing,
|
||||
'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
|
||||
with open(file_path, "w") as file:
|
||||
with open(file_path, 'w') as file:
|
||||
json.dump(variables, file)
|
||||
|
||||
return file_path
|
||||
@ -152,10 +157,10 @@ def open_configuration(
|
||||
original_file_path = file_path
|
||||
file_path = get_file_path(file_path)
|
||||
|
||||
if file_path != "" and file_path != None:
|
||||
if file_path != '' and file_path != None:
|
||||
print(file_path)
|
||||
# load variables from JSON file
|
||||
with open(file_path, "r") as f:
|
||||
with open(file_path, 'r') as f:
|
||||
my_data = json.load(f)
|
||||
else:
|
||||
file_path = original_file_path # In case a file_path was provided and the user decide to cancel the open action
|
||||
@ -164,36 +169,40 @@ def open_configuration(
|
||||
# Return the values of the variables as a dictionary
|
||||
return (
|
||||
file_path,
|
||||
my_data.get("pretrained_model_name_or_path", pretrained_model_name_or_path),
|
||||
my_data.get("v2", v2),
|
||||
my_data.get("v_parameterization", v_parameterization),
|
||||
my_data.get("logging_dir", logging_dir),
|
||||
my_data.get("train_data_dir", train_data_dir),
|
||||
my_data.get("reg_data_dir", reg_data_dir),
|
||||
my_data.get("output_dir", output_dir),
|
||||
my_data.get("max_resolution", max_resolution),
|
||||
my_data.get("learning_rate", learning_rate),
|
||||
my_data.get("lr_scheduler", lr_scheduler),
|
||||
my_data.get("lr_warmup", lr_warmup),
|
||||
my_data.get("train_batch_size", train_batch_size),
|
||||
my_data.get("epoch", epoch),
|
||||
my_data.get("save_every_n_epochs", save_every_n_epochs),
|
||||
my_data.get("mixed_precision", mixed_precision),
|
||||
my_data.get("save_precision", save_precision),
|
||||
my_data.get("seed", seed),
|
||||
my_data.get("num_cpu_threads_per_process", num_cpu_threads_per_process),
|
||||
my_data.get("convert_to_safetensors", convert_to_safetensors),
|
||||
my_data.get("convert_to_ckpt", convert_to_ckpt),
|
||||
my_data.get("cache_latent", cache_latent),
|
||||
my_data.get("caption_extention", caption_extention),
|
||||
my_data.get("use_safetensors", use_safetensors),
|
||||
my_data.get("enable_bucket", enable_bucket),
|
||||
my_data.get("gradient_checkpointing", gradient_checkpointing),
|
||||
my_data.get("full_fp16", full_fp16),
|
||||
my_data.get("no_token_padding", no_token_padding),
|
||||
my_data.get("stop_text_encoder_training", stop_text_encoder_training),
|
||||
my_data.get("use_8bit_adam", use_8bit_adam),
|
||||
my_data.get("xformers", xformers),
|
||||
my_data.get(
|
||||
'pretrained_model_name_or_path', pretrained_model_name_or_path
|
||||
),
|
||||
my_data.get('v2', v2),
|
||||
my_data.get('v_parameterization', v_parameterization),
|
||||
my_data.get('logging_dir', logging_dir),
|
||||
my_data.get('train_data_dir', train_data_dir),
|
||||
my_data.get('reg_data_dir', reg_data_dir),
|
||||
my_data.get('output_dir', output_dir),
|
||||
my_data.get('max_resolution', max_resolution),
|
||||
my_data.get('learning_rate', learning_rate),
|
||||
my_data.get('lr_scheduler', lr_scheduler),
|
||||
my_data.get('lr_warmup', lr_warmup),
|
||||
my_data.get('train_batch_size', train_batch_size),
|
||||
my_data.get('epoch', epoch),
|
||||
my_data.get('save_every_n_epochs', save_every_n_epochs),
|
||||
my_data.get('mixed_precision', mixed_precision),
|
||||
my_data.get('save_precision', save_precision),
|
||||
my_data.get('seed', seed),
|
||||
my_data.get(
|
||||
'num_cpu_threads_per_process', num_cpu_threads_per_process
|
||||
),
|
||||
my_data.get('convert_to_safetensors', convert_to_safetensors),
|
||||
my_data.get('convert_to_ckpt', convert_to_ckpt),
|
||||
my_data.get('cache_latent', cache_latent),
|
||||
my_data.get('caption_extention', caption_extention),
|
||||
my_data.get('use_safetensors', use_safetensors),
|
||||
my_data.get('enable_bucket', enable_bucket),
|
||||
my_data.get('gradient_checkpointing', gradient_checkpointing),
|
||||
my_data.get('full_fp16', full_fp16),
|
||||
my_data.get('no_token_padding', no_token_padding),
|
||||
my_data.get('stop_text_encoder_training', stop_text_encoder_training),
|
||||
my_data.get('use_8bit_adam', use_8bit_adam),
|
||||
my_data.get('xformers', xformers),
|
||||
)
|
||||
|
||||
|
||||
@ -229,46 +238,46 @@ def train_model(
|
||||
use_8bit_adam,
|
||||
xformers,
|
||||
):
|
||||
|
||||
def save_inference_file(output_dir, v2, v_parameterization):
|
||||
# Copy inference model for v2 if required
|
||||
if v2 and v_parameterization:
|
||||
print(f"Saving v2-inference-v.yaml as {output_dir}/last.yaml")
|
||||
print(f'Saving v2-inference-v.yaml as {output_dir}/last.yaml')
|
||||
shutil.copy(
|
||||
f"./v2_inference/v2-inference-v.yaml",
|
||||
f"{output_dir}/last.yaml",
|
||||
f'./v2_inference/v2-inference-v.yaml',
|
||||
f'{output_dir}/last.yaml',
|
||||
)
|
||||
elif v2:
|
||||
print(f"Saving v2-inference.yaml as {output_dir}/last.yaml")
|
||||
print(f'Saving v2-inference.yaml as {output_dir}/last.yaml')
|
||||
shutil.copy(
|
||||
f"./v2_inference/v2-inference.yaml",
|
||||
f"{output_dir}/last.yaml",
|
||||
f'./v2_inference/v2-inference.yaml',
|
||||
f'{output_dir}/last.yaml',
|
||||
)
|
||||
|
||||
if pretrained_model_name_or_path == "":
|
||||
msgbox("Source model information is missing")
|
||||
if pretrained_model_name_or_path == '':
|
||||
msgbox('Source model information is missing')
|
||||
return
|
||||
|
||||
if train_data_dir == "":
|
||||
msgbox("Image folder path is missing")
|
||||
if train_data_dir == '':
|
||||
msgbox('Image folder path is missing')
|
||||
return
|
||||
|
||||
if not os.path.exists(train_data_dir):
|
||||
msgbox("Image folder does not exist")
|
||||
msgbox('Image folder does not exist')
|
||||
return
|
||||
|
||||
if reg_data_dir != "":
|
||||
if reg_data_dir != '':
|
||||
if not os.path.exists(reg_data_dir):
|
||||
msgbox("Regularisation folder does not exist")
|
||||
msgbox('Regularisation folder does not exist')
|
||||
return
|
||||
|
||||
if output_dir == "":
|
||||
msgbox("Output folder path is missing")
|
||||
if output_dir == '':
|
||||
msgbox('Output folder path is missing')
|
||||
return
|
||||
|
||||
# Get a list of all subfolders in train_data_dir
|
||||
subfolders = [
|
||||
f for f in os.listdir(train_data_dir)
|
||||
f
|
||||
for f in os.listdir(train_data_dir)
|
||||
if os.path.isdir(os.path.join(train_data_dir, f))
|
||||
]
|
||||
|
||||
@ -277,115 +286,127 @@ def train_model(
|
||||
# Loop through each subfolder and extract the number of repeats
|
||||
for folder in subfolders:
|
||||
# Extract the number of repeats from the folder name
|
||||
repeats = int(folder.split("_")[0])
|
||||
repeats = int(folder.split('_')[0])
|
||||
|
||||
# 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(".jpg") or f.endswith(".jpeg") or f.endswith(".png")
|
||||
or f.endswith(".webp")
|
||||
])
|
||||
num_images = len(
|
||||
[
|
||||
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
|
||||
steps = repeats * num_images
|
||||
total_steps += steps
|
||||
|
||||
# Print the result
|
||||
print(f"Folder {folder}: {steps} steps")
|
||||
print(f'Folder {folder}: {steps} steps')
|
||||
|
||||
# Print the result
|
||||
# print(f"{total_steps} total steps")
|
||||
|
||||
if reg_data_dir == "":
|
||||
if reg_data_dir == '':
|
||||
reg_factor = 1
|
||||
else:
|
||||
print(
|
||||
"Regularisation images are used... Will double the number of steps required..."
|
||||
'Regularisation images are used... Will double the number of steps required...'
|
||||
)
|
||||
reg_factor = 2
|
||||
|
||||
# calculate max_train_steps
|
||||
max_train_steps = int(
|
||||
math.ceil(
|
||||
float(total_steps) / int(train_batch_size) * int(epoch) *
|
||||
int(reg_factor)))
|
||||
print(f"max_train_steps = {max_train_steps}")
|
||||
float(total_steps)
|
||||
/ int(train_batch_size)
|
||||
* int(epoch)
|
||||
* int(reg_factor)
|
||||
)
|
||||
)
|
||||
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}")
|
||||
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))
|
||||
print(f"lr_warmup_steps = {lr_warmup_steps}")
|
||||
print(f'lr_warmup_steps = {lr_warmup_steps}')
|
||||
|
||||
run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_db_fixed.py"'
|
||||
if v2:
|
||||
run_cmd += " --v2"
|
||||
run_cmd += ' --v2'
|
||||
if v_parameterization:
|
||||
run_cmd += " --v_parameterization"
|
||||
run_cmd += ' --v_parameterization'
|
||||
if cache_latent:
|
||||
run_cmd += " --cache_latents"
|
||||
run_cmd += ' --cache_latents'
|
||||
if use_safetensors:
|
||||
run_cmd += " --use_safetensors"
|
||||
run_cmd += ' --use_safetensors'
|
||||
if enable_bucket:
|
||||
run_cmd += " --enable_bucket"
|
||||
run_cmd += ' --enable_bucket'
|
||||
if gradient_checkpointing:
|
||||
run_cmd += " --gradient_checkpointing"
|
||||
run_cmd += ' --gradient_checkpointing'
|
||||
if full_fp16:
|
||||
run_cmd += " --full_fp16"
|
||||
run_cmd += ' --full_fp16'
|
||||
if no_token_padding:
|
||||
run_cmd += " --no_token_padding"
|
||||
run_cmd += ' --no_token_padding'
|
||||
if use_8bit_adam:
|
||||
run_cmd += " --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 += ' --xformers'
|
||||
run_cmd += (
|
||||
f' --pretrained_model_name_or_path={pretrained_model_name_or_path}'
|
||||
)
|
||||
run_cmd += f' --train_data_dir="{train_data_dir}"'
|
||||
if len(reg_data_dir):
|
||||
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}"
|
||||
run_cmd += f" --caption_extention={caption_extention}"
|
||||
run_cmd += f" --stop_text_encoder_training={stop_text_encoder_training}"
|
||||
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}'
|
||||
run_cmd += f' --caption_extention={caption_extention}'
|
||||
run_cmd += f' --stop_text_encoder_training={stop_text_encoder_training}'
|
||||
|
||||
print(run_cmd)
|
||||
# Run the command
|
||||
subprocess.run(run_cmd)
|
||||
|
||||
# check if output_dir/last is a directory... therefore it is a diffuser model
|
||||
last_dir = pathlib.Path(f"{output_dir}/last")
|
||||
last_dir = pathlib.Path(f'{output_dir}/last')
|
||||
print(last_dir)
|
||||
if last_dir.is_dir():
|
||||
if convert_to_ckpt:
|
||||
print(f"Converting diffuser model {last_dir} to {last_dir}.ckpt")
|
||||
print(f'Converting diffuser model {last_dir} to {last_dir}.ckpt')
|
||||
os.system(
|
||||
f"python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.ckpt --{save_precision}"
|
||||
f'python ./tools/convert_diffusers20_original_sd.py {last_dir} {last_dir}.ckpt --{save_precision}'
|
||||
)
|
||||
|
||||
save_inference_file(output_dir, v2, v_parameterization)
|
||||
|
||||
if convert_to_safetensors:
|
||||
print(
|
||||
f"Converting diffuser model {last_dir} to {last_dir}.safetensors"
|
||||
f'Converting diffuser model {last_dir} to {last_dir}.safetensors'
|
||||
)
|
||||
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}'
|
||||
)
|
||||
|
||||
save_inference_file(output_dir, v2, v_parameterization)
|
||||
@ -400,13 +421,13 @@ def train_model(
|
||||
def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
|
||||
# define a list of substrings to search for
|
||||
substrings_v2 = [
|
||||
"stabilityai/stable-diffusion-2-1-base",
|
||||
"stabilityai/stable-diffusion-2-base",
|
||||
'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
|
||||
if str(value) in substrings_v2:
|
||||
print("SD v2 model detected. Setting --v2 parameter")
|
||||
print('SD v2 model detected. Setting --v2 parameter')
|
||||
v2 = True
|
||||
v_parameterization = False
|
||||
|
||||
@ -414,14 +435,14 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
|
||||
|
||||
# define a list of substrings to search for v-objective
|
||||
substrings_v_parameterization = [
|
||||
"stabilityai/stable-diffusion-2-1",
|
||||
"stabilityai/stable-diffusion-2",
|
||||
'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
|
||||
if str(value) in substrings_v_parameterization:
|
||||
print(
|
||||
"SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization"
|
||||
'SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization'
|
||||
)
|
||||
v2 = True
|
||||
v_parameterization = True
|
||||
@ -430,8 +451,8 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
|
||||
|
||||
# define a list of substrings to v1.x
|
||||
substrings_v1_model = [
|
||||
"CompVis/stable-diffusion-v1-4",
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
'CompVis/stable-diffusion-v1-4',
|
||||
'runwayml/stable-diffusion-v1-5',
|
||||
]
|
||||
|
||||
if str(value) in substrings_v1_model:
|
||||
@ -440,62 +461,67 @@ def set_pretrained_model_name_or_path_input(value, v2, v_parameterization):
|
||||
|
||||
return value, v2, v_parameterization
|
||||
|
||||
if value == "custom":
|
||||
value = ""
|
||||
if value == 'custom':
|
||||
value = ''
|
||||
v2 = False
|
||||
v_parameterization = False
|
||||
|
||||
return value, v2, v_parameterization
|
||||
|
||||
|
||||
css = ""
|
||||
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"
|
||||
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:
|
||||
dummy_true = gr.Label(value=True, visible=False)
|
||||
dummy_false = gr.Label(value=False, visible=False)
|
||||
gr.Markdown("Enter kohya finetuner parameter using this interface.")
|
||||
with gr.Accordion("Configuration File Load/Save", open=False):
|
||||
gr.Markdown('Enter kohya finetuner parameter using this interface.')
|
||||
with gr.Accordion('Configuration File Load/Save', 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")
|
||||
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...")
|
||||
config_file_name.change(remove_doublequote,
|
||||
inputs=[config_file_name],
|
||||
outputs=[config_file_name])
|
||||
with gr.Tab("Source model"):
|
||||
label='',
|
||||
placeholder="type the configuration file path or use the 'Open' button above to select it...",
|
||||
)
|
||||
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",
|
||||
label='Pretrained model name or path',
|
||||
placeholder='enter the path to custom model or name of pretrained model',
|
||||
)
|
||||
model_list = gr.Dropdown(
|
||||
label="(Optional) Model Quick Pick",
|
||||
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",
|
||||
'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',
|
||||
],
|
||||
)
|
||||
with gr.Row():
|
||||
v2_input = gr.Checkbox(label="v2", value=True)
|
||||
v_parameterization_input = gr.Checkbox(label="v_parameterization",
|
||||
value=False)
|
||||
v2_input = gr.Checkbox(label='v2', value=True)
|
||||
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],
|
||||
@ -511,44 +537,49 @@ with interface:
|
||||
],
|
||||
)
|
||||
|
||||
with gr.Tab("Directories"):
|
||||
with gr.Tab('Directories'):
|
||||
with gr.Row():
|
||||
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',
|
||||
)
|
||||
train_data_dir_input_folder = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
train_data_dir_input_folder.click(get_folder_path,
|
||||
outputs=train_data_dir_input)
|
||||
reg_data_dir_input = gr.Textbox(
|
||||
label="Regularisation folder",
|
||||
placeholder=
|
||||
"(Optional) Directory where where the regularization folders containing the images are located",
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
train_data_dir_input_folder.click(
|
||||
get_folder_path, outputs=train_data_dir_input
|
||||
)
|
||||
reg_data_dir_input = gr.Textbox(
|
||||
label='Regularisation folder',
|
||||
placeholder='(Optional) Directory 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(
|
||||
get_folder_path, outputs=reg_data_dir_input
|
||||
)
|
||||
reg_data_dir_input_folder = gr.Button("📂",
|
||||
elem_id="open_folder_small")
|
||||
reg_data_dir_input_folder.click(get_folder_path,
|
||||
outputs=reg_data_dir_input)
|
||||
with gr.Row():
|
||||
output_dir_input = gr.Textbox(
|
||||
label="Output directory",
|
||||
placeholder="Directory to output trained model",
|
||||
label='Output directory',
|
||||
placeholder='Directory 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_input
|
||||
)
|
||||
output_dir_input_folder = gr.Button("📂",
|
||||
elem_id="open_folder_small")
|
||||
output_dir_input_folder.click(get_folder_path,
|
||||
outputs=output_dir_input)
|
||||
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',
|
||||
)
|
||||
logging_dir_input_folder = gr.Button(
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
logging_dir_input_folder.click(
|
||||
get_folder_path, outputs=logging_dir_input
|
||||
)
|
||||
logging_dir_input_folder = gr.Button("📂",
|
||||
elem_id="open_folder_small")
|
||||
logging_dir_input_folder.click(get_folder_path,
|
||||
outputs=logging_dir_input)
|
||||
train_data_dir_input.change(
|
||||
remove_doublequote,
|
||||
inputs=[train_data_dir_input],
|
||||
@ -559,111 +590,130 @@ with interface:
|
||||
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"):
|
||||
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.Row():
|
||||
learning_rate_input = gr.Textbox(label="Learning rate", value=1e-6)
|
||||
learning_rate_input = gr.Textbox(label='Learning rate', value=1e-6)
|
||||
lr_scheduler_input = gr.Dropdown(
|
||||
label="LR Scheduler",
|
||||
label='LR Scheduler',
|
||||
choices=[
|
||||
"constant",
|
||||
"constant_with_warmup",
|
||||
"cosine",
|
||||
"cosine_with_restarts",
|
||||
"linear",
|
||||
"polynomial",
|
||||
'constant',
|
||||
'constant_with_warmup',
|
||||
'cosine',
|
||||
'cosine_with_restarts',
|
||||
'linear',
|
||||
'polynomial',
|
||||
],
|
||||
value="constant",
|
||||
value='constant',
|
||||
)
|
||||
lr_warmup_input = gr.Textbox(label="LR warmup", value=0)
|
||||
lr_warmup_input = gr.Textbox(label='LR warmup', value=0)
|
||||
with gr.Row():
|
||||
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)
|
||||
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",
|
||||
label='Mixed precision',
|
||||
choices=[
|
||||
"no",
|
||||
"fp16",
|
||||
"bf16",
|
||||
'no',
|
||||
'fp16',
|
||||
'bf16',
|
||||
],
|
||||
value="fp16",
|
||||
value='fp16',
|
||||
)
|
||||
save_precision_input = gr.Dropdown(
|
||||
label="Save precision",
|
||||
label='Save precision',
|
||||
choices=[
|
||||
"float",
|
||||
"fp16",
|
||||
"bf16",
|
||||
'float',
|
||||
'fp16',
|
||||
'bf16',
|
||||
],
|
||||
value="fp16",
|
||||
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",
|
||||
label='Number of CPU threads per process',
|
||||
value=os.cpu_count(),
|
||||
)
|
||||
with gr.Row():
|
||||
seed_input = gr.Textbox(label="Seed", value=1234)
|
||||
max_resolution_input = gr.Textbox(label="Max resolution",
|
||||
value="512,512",
|
||||
placeholder="512,512")
|
||||
seed_input = gr.Textbox(label='Seed', value=1234)
|
||||
max_resolution_input = gr.Textbox(
|
||||
label='Max resolution', value='512,512', placeholder='512,512'
|
||||
)
|
||||
with gr.Row():
|
||||
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",
|
||||
label='Stop text encoder training',
|
||||
)
|
||||
with gr.Row():
|
||||
full_fp16_input = gr.Checkbox(
|
||||
label="Full fp16 training (experimental)", value=False)
|
||||
no_token_padding_input = gr.Checkbox(label="No token padding",
|
||||
value=False)
|
||||
label='Full fp16 training (experimental)', value=False
|
||||
)
|
||||
no_token_padding_input = gr.Checkbox(
|
||||
label='No token padding', value=False
|
||||
)
|
||||
use_safetensors_input = gr.Checkbox(
|
||||
label="Use safetensor when saving", value=False)
|
||||
label='Use safetensor when saving', value=False
|
||||
)
|
||||
|
||||
gradient_checkpointing_input = gr.Checkbox(
|
||||
label="Gradient checkpointing", value=False)
|
||||
label='Gradient checkpointing', value=False
|
||||
)
|
||||
with gr.Row():
|
||||
enable_bucket_input = gr.Checkbox(label="Enable buckets",
|
||||
value=True)
|
||||
cache_latent_input = gr.Checkbox(label="Cache latent", value=True)
|
||||
use_8bit_adam_input = gr.Checkbox(label="Use 8bit adam",
|
||||
value=True)
|
||||
xformers_input = gr.Checkbox(label="Use xformers", value=True)
|
||||
enable_bucket_input = gr.Checkbox(
|
||||
label='Enable buckets', value=True
|
||||
)
|
||||
cache_latent_input = gr.Checkbox(label='Cache latent', value=True)
|
||||
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(
|
||||
label="Convert to SafeTensors", value=True)
|
||||
convert_to_ckpt_input = gr.Checkbox(label="Convert to CKPT",
|
||||
value=False)
|
||||
with gr.Tab("Utilities"):
|
||||
label='Convert to SafeTensors', value=True
|
||||
)
|
||||
convert_to_ckpt_input = gr.Checkbox(
|
||||
label='Convert to CKPT', value=False
|
||||
)
|
||||
with gr.Tab('Utilities'):
|
||||
# Dreambooth folder creation tab
|
||||
gradio_dreambooth_folder_creation_tab(train_data_dir_input, reg_data_dir_input, output_dir_input, logging_dir_input)
|
||||
gradio_dreambooth_folder_creation_tab(
|
||||
train_data_dir_input,
|
||||
reg_data_dir_input,
|
||||
output_dir_input,
|
||||
logging_dir_input,
|
||||
)
|
||||
# Captionning tab
|
||||
gradio_caption_gui_tab()
|
||||
|
||||
button_run = gr.Button("Train model")
|
||||
button_run = gr.Button('Train model')
|
||||
|
||||
button_open_config.click(
|
||||
open_configuration,
|
||||
|
||||
@ -3,69 +3,85 @@ from easygui import msgbox
|
||||
import subprocess
|
||||
from .common_gui import get_folder_path
|
||||
|
||||
def caption_images(caption_text_input, images_dir_input, overwrite_input, caption_file_ext):
|
||||
|
||||
def caption_images(
|
||||
caption_text_input, images_dir_input, overwrite_input, caption_file_ext
|
||||
):
|
||||
# Check for caption_text_input
|
||||
if caption_text_input == "":
|
||||
msgbox("Caption text is missing...")
|
||||
if caption_text_input == '':
|
||||
msgbox('Caption text is missing...')
|
||||
return
|
||||
|
||||
|
||||
# Check for images_dir_input
|
||||
if images_dir_input == "":
|
||||
msgbox("Image folder is missing...")
|
||||
if images_dir_input == '':
|
||||
msgbox('Image folder is missing...')
|
||||
return
|
||||
|
||||
print(f"Captionning files in {images_dir_input} with {caption_text_input}...")
|
||||
|
||||
print(
|
||||
f'Captionning files in {images_dir_input} with {caption_text_input}...'
|
||||
)
|
||||
run_cmd = f'python "tools/caption.py"'
|
||||
run_cmd += f' --caption_text="{caption_text_input}"'
|
||||
if overwrite_input:
|
||||
run_cmd += f' --overwrite'
|
||||
if caption_file_ext != "":
|
||||
if caption_file_ext != '':
|
||||
run_cmd += f' --caption_file_ext="{caption_file_ext}"'
|
||||
run_cmd += f' "{images_dir_input}"'
|
||||
|
||||
|
||||
print(run_cmd)
|
||||
|
||||
|
||||
# Run the command
|
||||
subprocess.run(run_cmd)
|
||||
|
||||
print("...captionning done")
|
||||
|
||||
print('...captionning done')
|
||||
|
||||
|
||||
###
|
||||
# Gradio UI
|
||||
###
|
||||
|
||||
|
||||
def gradio_caption_gui_tab():
|
||||
with gr.Tab("Captionning"):
|
||||
with gr.Tab('Captionning'):
|
||||
gr.Markdown(
|
||||
"This utility will allow the creation of caption files for each images in a folder."
|
||||
'This utility will allow the creation of caption files for each images in a folder.'
|
||||
)
|
||||
with gr.Row():
|
||||
caption_text_input = gr.Textbox(
|
||||
label="Caption text",
|
||||
placeholder="Eg: , by some artist",
|
||||
label='Caption text',
|
||||
placeholder='Eg: , by some artist',
|
||||
interactive=True,
|
||||
)
|
||||
)
|
||||
overwrite_input = gr.Checkbox(
|
||||
label="Overwrite existing captions in folder",
|
||||
label='Overwrite existing captions in folder',
|
||||
interactive=True,
|
||||
value=False
|
||||
value=False,
|
||||
)
|
||||
caption_file_ext = gr.Textbox(
|
||||
label="Caption file extension",
|
||||
placeholder="(Optional) Default: .caption",
|
||||
label='Caption file extension',
|
||||
placeholder='(Optional) Default: .caption',
|
||||
interactive=True,
|
||||
)
|
||||
with gr.Row():
|
||||
images_dir_input = gr.Textbox(
|
||||
label="Image forder to caption",
|
||||
placeholder="Directory containing the images to caption",
|
||||
label='Image forder to caption',
|
||||
placeholder='Directory containing the images to caption',
|
||||
interactive=True,
|
||||
)
|
||||
button_images_dir_input = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
button_images_dir_input.click(
|
||||
get_folder_path, outputs=images_dir_input)
|
||||
caption_button = gr.Button("Caption images")
|
||||
|
||||
caption_button.click(caption_images, inputs=[caption_text_input, images_dir_input, overwrite_input, caption_file_ext])
|
||||
|
||||
get_folder_path, outputs=images_dir_input
|
||||
)
|
||||
caption_button = gr.Button('Caption images')
|
||||
|
||||
caption_button.click(
|
||||
caption_images,
|
||||
inputs=[
|
||||
caption_text_input,
|
||||
images_dir_input,
|
||||
overwrite_input,
|
||||
caption_file_ext,
|
||||
],
|
||||
)
|
||||
|
||||
@ -1,19 +1,22 @@
|
||||
from easygui import diropenbox, fileopenbox
|
||||
|
||||
|
||||
def get_folder_path():
|
||||
folder_path = diropenbox("Select the directory to use")
|
||||
folder_path = diropenbox('Select the directory to use')
|
||||
|
||||
return folder_path
|
||||
|
||||
|
||||
def remove_doublequote(file_path):
|
||||
if file_path != None:
|
||||
file_path = file_path.replace('"', "")
|
||||
file_path = file_path.replace('"', '')
|
||||
|
||||
return file_path
|
||||
|
||||
def get_file_path(file_path):
|
||||
file_path = fileopenbox("Select the config file to load",
|
||||
default=file_path,
|
||||
filetypes="*.json")
|
||||
|
||||
return file_path
|
||||
def get_file_path(file_path):
|
||||
file_path = fileopenbox(
|
||||
'Select the config file to load', default=file_path, filetypes='*.json'
|
||||
)
|
||||
|
||||
return file_path
|
||||
|
||||
@ -4,14 +4,15 @@ from .common_gui import get_folder_path
|
||||
import shutil
|
||||
import os
|
||||
|
||||
|
||||
def copy_info_to_Directories_tab(training_folder):
|
||||
img_folder = os.path.join(training_folder, "img")
|
||||
if os.path.exists(os.path.join(training_folder, "reg")):
|
||||
reg_folder = os.path.join(training_folder, "reg")
|
||||
img_folder = os.path.join(training_folder, 'img')
|
||||
if os.path.exists(os.path.join(training_folder, 'reg')):
|
||||
reg_folder = os.path.join(training_folder, 'reg')
|
||||
else:
|
||||
reg_folder = ""
|
||||
model_folder = os.path.join(training_folder, "model")
|
||||
log_folder = os.path.join(training_folder, "log")
|
||||
reg_folder = ''
|
||||
model_folder = os.path.join(training_folder, 'model')
|
||||
log_folder = os.path.join(training_folder, 'log')
|
||||
|
||||
return img_folder, reg_folder, model_folder, log_folder
|
||||
|
||||
@ -27,7 +28,7 @@ def dreambooth_folder_preparation(
|
||||
):
|
||||
|
||||
# Check if the input variables are empty
|
||||
if (not len(util_training_dir_output)):
|
||||
if not len(util_training_dir_output):
|
||||
print(
|
||||
"Destination training directory is missing... can't perform the required task..."
|
||||
)
|
||||
@ -37,17 +38,17 @@ def dreambooth_folder_preparation(
|
||||
os.makedirs(util_training_dir_output, exist_ok=True)
|
||||
|
||||
# Check for instance prompt
|
||||
if util_instance_prompt_input == "":
|
||||
msgbox("Instance prompt missing...")
|
||||
if util_instance_prompt_input == '':
|
||||
msgbox('Instance prompt missing...')
|
||||
return
|
||||
|
||||
|
||||
# Check for class prompt
|
||||
if util_class_prompt_input == "":
|
||||
msgbox("Class prompt missing...")
|
||||
if util_class_prompt_input == '':
|
||||
msgbox('Class prompt missing...')
|
||||
return
|
||||
|
||||
# Create the training_dir path
|
||||
if (util_training_images_dir_input == ""):
|
||||
if util_training_images_dir_input == '':
|
||||
print(
|
||||
"Training images directory is missing... can't perform the required task..."
|
||||
)
|
||||
@ -55,106 +56,120 @@ def dreambooth_folder_preparation(
|
||||
else:
|
||||
training_dir = os.path.join(
|
||||
util_training_dir_output,
|
||||
f"img/{int(util_training_images_repeat_input)}_{util_instance_prompt_input} {util_class_prompt_input}",
|
||||
f'img/{int(util_training_images_repeat_input)}_{util_instance_prompt_input} {util_class_prompt_input}',
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(training_dir):
|
||||
print(f"Removing existing directory {training_dir}...")
|
||||
print(f'Removing existing directory {training_dir}...')
|
||||
shutil.rmtree(training_dir)
|
||||
|
||||
# Copy the training images to their respective directories
|
||||
print(f"Copy {util_training_images_dir_input} to {training_dir}...")
|
||||
print(f'Copy {util_training_images_dir_input} to {training_dir}...')
|
||||
shutil.copytree(util_training_images_dir_input, training_dir)
|
||||
|
||||
# Create the regularization_dir path
|
||||
if (not (util_class_prompt_input == "")
|
||||
or not util_regularization_images_repeat_input > 0):
|
||||
if (
|
||||
not (util_class_prompt_input == '')
|
||||
or not util_regularization_images_repeat_input > 0
|
||||
):
|
||||
print(
|
||||
"Regularization images directory or repeats is missing... not copying regularisation images..."
|
||||
'Regularization images directory or repeats is missing... not copying regularisation images...'
|
||||
)
|
||||
else:
|
||||
regularization_dir = os.path.join(
|
||||
util_training_dir_output,
|
||||
f"reg/{int(util_regularization_images_repeat_input)}_{util_class_prompt_input}",
|
||||
f'reg/{int(util_regularization_images_repeat_input)}_{util_class_prompt_input}',
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(regularization_dir):
|
||||
print(f"Removing existing directory {regularization_dir}...")
|
||||
print(f'Removing existing directory {regularization_dir}...')
|
||||
shutil.rmtree(regularization_dir)
|
||||
|
||||
# Copy the regularisation images to their respective directories
|
||||
print(
|
||||
f"Copy {util_regularization_images_dir_input} to {regularization_dir}..."
|
||||
f'Copy {util_regularization_images_dir_input} to {regularization_dir}...'
|
||||
)
|
||||
shutil.copytree(
|
||||
util_regularization_images_dir_input, regularization_dir
|
||||
)
|
||||
shutil.copytree(util_regularization_images_dir_input,
|
||||
regularization_dir)
|
||||
|
||||
print(
|
||||
f"Done creating kohya_ss training folder structure at {util_training_dir_output}..."
|
||||
f'Done creating kohya_ss training folder structure at {util_training_dir_output}...'
|
||||
)
|
||||
|
||||
def gradio_dreambooth_folder_creation_tab(train_data_dir_input, reg_data_dir_input, output_dir_input, logging_dir_input):
|
||||
with gr.Tab("Dreambooth folder preparation"):
|
||||
|
||||
def gradio_dreambooth_folder_creation_tab(
|
||||
train_data_dir_input,
|
||||
reg_data_dir_input,
|
||||
output_dir_input,
|
||||
logging_dir_input,
|
||||
):
|
||||
with gr.Tab('Dreambooth folder preparation'):
|
||||
gr.Markdown(
|
||||
"This utility will create the necessary folder structure for the training images and optional regularization images needed for the kohys_ss Dreambooth method to function correctly."
|
||||
'This utility will create the necessary folder structure for the training images and optional regularization images needed for the kohys_ss Dreambooth method to function correctly.'
|
||||
)
|
||||
with gr.Row():
|
||||
util_instance_prompt_input = gr.Textbox(
|
||||
label="Instance prompt",
|
||||
placeholder="Eg: asd",
|
||||
label='Instance prompt',
|
||||
placeholder='Eg: asd',
|
||||
interactive=True,
|
||||
)
|
||||
util_class_prompt_input = gr.Textbox(
|
||||
label="Class prompt",
|
||||
placeholder="Eg: person",
|
||||
label='Class prompt',
|
||||
placeholder='Eg: person',
|
||||
interactive=True,
|
||||
)
|
||||
with gr.Row():
|
||||
util_training_images_dir_input = gr.Textbox(
|
||||
label="Training images",
|
||||
placeholder="Directory containing the training images",
|
||||
label='Training images',
|
||||
placeholder='Directory containing the training images',
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_images_dir_input = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
button_util_training_images_dir_input.click(
|
||||
get_folder_path, outputs=util_training_images_dir_input)
|
||||
get_folder_path, outputs=util_training_images_dir_input
|
||||
)
|
||||
util_training_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
label='Repeats',
|
||||
value=40,
|
||||
interactive=True,
|
||||
elem_id="number_input")
|
||||
elem_id='number_input',
|
||||
)
|
||||
with gr.Row():
|
||||
util_regularization_images_dir_input = gr.Textbox(
|
||||
label="Regularisation images",
|
||||
placeholder=
|
||||
"(Optional) Directory containing the regularisation images",
|
||||
label='Regularisation images',
|
||||
placeholder='(Optional) Directory containing the regularisation images',
|
||||
interactive=True,
|
||||
)
|
||||
button_util_regularization_images_dir_input = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
button_util_regularization_images_dir_input.click(
|
||||
get_folder_path,
|
||||
outputs=util_regularization_images_dir_input)
|
||||
get_folder_path, outputs=util_regularization_images_dir_input
|
||||
)
|
||||
util_regularization_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
label='Repeats',
|
||||
value=1,
|
||||
interactive=True,
|
||||
elem_id="number_input")
|
||||
elem_id='number_input',
|
||||
)
|
||||
with gr.Row():
|
||||
util_training_dir_output = gr.Textbox(
|
||||
label="Destination training directory",
|
||||
placeholder=
|
||||
"Directory where formatted training and regularisation folders will be placed",
|
||||
label='Destination training directory',
|
||||
placeholder='Directory where formatted training and regularisation folders will be placed',
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_dir_output = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
'📂', elem_id='open_folder_small'
|
||||
)
|
||||
button_util_training_dir_output.click(
|
||||
get_folder_path, outputs=util_training_dir_output)
|
||||
button_prepare_training_data = gr.Button("Prepare training data")
|
||||
get_folder_path, outputs=util_training_dir_output
|
||||
)
|
||||
button_prepare_training_data = gr.Button('Prepare training data')
|
||||
button_prepare_training_data.click(
|
||||
dreambooth_folder_preparation,
|
||||
inputs=[
|
||||
@ -168,12 +183,15 @@ def gradio_dreambooth_folder_creation_tab(train_data_dir_input, reg_data_dir_inp
|
||||
],
|
||||
)
|
||||
button_copy_info_to_Directories_tab = gr.Button(
|
||||
"Copy info to Directories Tab")
|
||||
button_copy_info_to_Directories_tab.click(copy_info_to_Directories_tab,
|
||||
inputs=[util_training_dir_output],
|
||||
outputs=[
|
||||
train_data_dir_input,
|
||||
reg_data_dir_input,
|
||||
output_dir_input,
|
||||
logging_dir_input
|
||||
])
|
||||
'Copy info to Directories Tab'
|
||||
)
|
||||
button_copy_info_to_Directories_tab.click(
|
||||
copy_info_to_Directories_tab,
|
||||
inputs=[util_training_dir_output],
|
||||
outputs=[
|
||||
train_data_dir_input,
|
||||
reg_data_dir_input,
|
||||
output_dir_input,
|
||||
logging_dir_input,
|
||||
],
|
||||
)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user