Refactor GUI code to be modular
This commit is contained in:
bmaltais
parent
9bddf2ce4c
commit
90dad5471c
@ -10,11 +10,12 @@ import os
|
||||
import subprocess
|
||||
import pathlib
|
||||
import shutil
|
||||
import dreambooth_gui.caption_gui as caption_gui
|
||||
from glob import glob
|
||||
from os.path import join
|
||||
from easygui import fileopenbox, filesavebox, enterbox, diropenbox, msgbox
|
||||
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 easygui import filesavebox, msgbox
|
||||
|
||||
# sys.path.insert(0, './dreambooth_gui')
|
||||
|
||||
def save_configuration(
|
||||
save_as,
|
||||
@ -163,8 +164,7 @@ 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("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),
|
||||
@ -181,8 +181,7 @@ def open_configuration(
|
||||
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("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),
|
||||
@ -449,119 +448,6 @@ def set_pretrained_model_name_or_path_input(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
|
||||
|
||||
|
||||
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_folder_path():
|
||||
folder_path = diropenbox("Select the directory to use")
|
||||
|
||||
return folder_path
|
||||
|
||||
|
||||
def dreambooth_folder_preparation(
|
||||
util_training_images_dir_input,
|
||||
util_training_images_repeat_input,
|
||||
util_instance_prompt_input,
|
||||
util_regularization_images_dir_input,
|
||||
util_regularization_images_repeat_input,
|
||||
util_class_prompt_input,
|
||||
util_training_dir_output,
|
||||
):
|
||||
|
||||
# Check if the input variables are empty
|
||||
if (not len(util_training_dir_output)):
|
||||
print(
|
||||
"Destination training directory is missing... can't perform the required task..."
|
||||
)
|
||||
return
|
||||
else:
|
||||
# Create the util_training_dir_output directory if it doesn't exist
|
||||
os.makedirs(util_training_dir_output, exist_ok=True)
|
||||
|
||||
# Check for instance prompt
|
||||
if util_instance_prompt_input == "":
|
||||
msgbox("Instance prompt missing...")
|
||||
return
|
||||
|
||||
# Check for class prompt
|
||||
if util_class_prompt_input == "":
|
||||
msgbox("Class prompt missing...")
|
||||
return
|
||||
|
||||
# Create the training_dir path
|
||||
if (util_training_images_dir_input == ""):
|
||||
print(
|
||||
"Training images directory is missing... can't perform the required task..."
|
||||
)
|
||||
return
|
||||
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}",
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(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}...")
|
||||
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):
|
||||
print(
|
||||
"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}",
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(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}..."
|
||||
)
|
||||
shutil.copytree(util_regularization_images_dir_input,
|
||||
regularization_dir)
|
||||
|
||||
print(
|
||||
f"Done creating kohya_ss training folder structure at {util_training_dir_output}..."
|
||||
)
|
||||
|
||||
|
||||
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")
|
||||
else:
|
||||
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
|
||||
|
||||
|
||||
css = ""
|
||||
|
||||
if os.path.exists("./style.css"):
|
||||
@ -771,94 +657,14 @@ with interface:
|
||||
label="Convert to SafeTensors", value=True)
|
||||
convert_to_ckpt_input = gr.Checkbox(label="Convert to CKPT",
|
||||
value=False)
|
||||
|
||||
with gr.Tab("Utilities"):
|
||||
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."
|
||||
)
|
||||
with gr.Row():
|
||||
util_instance_prompt_input = gr.Textbox(
|
||||
label="Instance prompt",
|
||||
placeholder="Eg: asd",
|
||||
interactive=True,
|
||||
)
|
||||
util_class_prompt_input = gr.Textbox(
|
||||
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",
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_images_dir_input = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
button_util_training_images_dir_input.click(
|
||||
get_folder_path, outputs=util_training_images_dir_input)
|
||||
util_training_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
value=40,
|
||||
interactive=True,
|
||||
elem_id="number_input")
|
||||
with gr.Row():
|
||||
util_regularization_images_dir_input = gr.Textbox(
|
||||
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")
|
||||
button_util_regularization_images_dir_input.click(
|
||||
get_folder_path,
|
||||
outputs=util_regularization_images_dir_input)
|
||||
util_regularization_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
value=1,
|
||||
interactive=True,
|
||||
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",
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_dir_output = gr.Button(
|
||||
"📂", 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")
|
||||
button_prepare_training_data.click(
|
||||
dreambooth_folder_preparation,
|
||||
inputs=[
|
||||
util_training_images_dir_input,
|
||||
util_training_images_repeat_input,
|
||||
util_instance_prompt_input,
|
||||
util_regularization_images_dir_input,
|
||||
util_regularization_images_repeat_input,
|
||||
util_class_prompt_input,
|
||||
util_training_dir_output,
|
||||
],
|
||||
)
|
||||
button_copy_info_to_Directories_tab = gr.Button(
|
||||
"Copy info to Directories Tab")
|
||||
caption_gui.gradio_caption_gui()
|
||||
# Dreambooth folder creation tab
|
||||
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_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
|
||||
])
|
||||
|
||||
button_open_config.click(
|
||||
open_configuration,
|
||||
inputs=[
|
||||
|
||||
@ -1,11 +1,7 @@
|
||||
import gradio as gr
|
||||
from easygui import diropenbox, msgbox
|
||||
from easygui import msgbox
|
||||
import subprocess
|
||||
|
||||
def get_folder_path():
|
||||
folder_path = diropenbox("Select the directory to use")
|
||||
|
||||
return folder_path
|
||||
from .common_gui import get_folder_path
|
||||
|
||||
def caption_images(caption_text_input, images_dir_input, overwrite_input, caption_file_ext):
|
||||
# Check for caption_text_input
|
||||
@ -38,7 +34,7 @@ def caption_images(caption_text_input, images_dir_input, overwrite_input, captio
|
||||
# Gradio UI
|
||||
###
|
||||
|
||||
def gradio_caption_gui():
|
||||
def gradio_caption_gui_tab():
|
||||
with gr.Tab("Captionning"):
|
||||
gr.Markdown(
|
||||
"This utility will allow the creation of caption files for each images in a folder."
|
||||
|
||||
19
dreambooth_gui/common_gui.py
Normal file
19
dreambooth_gui/common_gui.py
Normal file
@ -0,0 +1,19 @@
|
||||
from easygui import diropenbox, fileopenbox
|
||||
|
||||
def get_folder_path():
|
||||
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('"', "")
|
||||
|
||||
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
|
||||
179
dreambooth_gui/dreambooth_folder_creation.py
Normal file
179
dreambooth_gui/dreambooth_folder_creation.py
Normal file
@ -0,0 +1,179 @@
|
||||
import gradio as gr
|
||||
from easygui import diropenbox, msgbox
|
||||
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")
|
||||
else:
|
||||
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
|
||||
|
||||
|
||||
def dreambooth_folder_preparation(
|
||||
util_training_images_dir_input,
|
||||
util_training_images_repeat_input,
|
||||
util_instance_prompt_input,
|
||||
util_regularization_images_dir_input,
|
||||
util_regularization_images_repeat_input,
|
||||
util_class_prompt_input,
|
||||
util_training_dir_output,
|
||||
):
|
||||
|
||||
# Check if the input variables are empty
|
||||
if (not len(util_training_dir_output)):
|
||||
print(
|
||||
"Destination training directory is missing... can't perform the required task..."
|
||||
)
|
||||
return
|
||||
else:
|
||||
# Create the util_training_dir_output directory if it doesn't exist
|
||||
os.makedirs(util_training_dir_output, exist_ok=True)
|
||||
|
||||
# Check for instance prompt
|
||||
if util_instance_prompt_input == "":
|
||||
msgbox("Instance prompt missing...")
|
||||
return
|
||||
|
||||
# Check for class prompt
|
||||
if util_class_prompt_input == "":
|
||||
msgbox("Class prompt missing...")
|
||||
return
|
||||
|
||||
# Create the training_dir path
|
||||
if (util_training_images_dir_input == ""):
|
||||
print(
|
||||
"Training images directory is missing... can't perform the required task..."
|
||||
)
|
||||
return
|
||||
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}",
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(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}...")
|
||||
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):
|
||||
print(
|
||||
"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}",
|
||||
)
|
||||
|
||||
# Remove folders if they exist
|
||||
if os.path.exists(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}..."
|
||||
)
|
||||
shutil.copytree(util_regularization_images_dir_input,
|
||||
regularization_dir)
|
||||
|
||||
print(
|
||||
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"):
|
||||
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."
|
||||
)
|
||||
with gr.Row():
|
||||
util_instance_prompt_input = gr.Textbox(
|
||||
label="Instance prompt",
|
||||
placeholder="Eg: asd",
|
||||
interactive=True,
|
||||
)
|
||||
util_class_prompt_input = gr.Textbox(
|
||||
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",
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_images_dir_input = gr.Button(
|
||||
"📂", elem_id="open_folder_small")
|
||||
button_util_training_images_dir_input.click(
|
||||
get_folder_path, outputs=util_training_images_dir_input)
|
||||
util_training_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
value=40,
|
||||
interactive=True,
|
||||
elem_id="number_input")
|
||||
with gr.Row():
|
||||
util_regularization_images_dir_input = gr.Textbox(
|
||||
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")
|
||||
button_util_regularization_images_dir_input.click(
|
||||
get_folder_path,
|
||||
outputs=util_regularization_images_dir_input)
|
||||
util_regularization_images_repeat_input = gr.Number(
|
||||
label="Repeats",
|
||||
value=1,
|
||||
interactive=True,
|
||||
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",
|
||||
interactive=True,
|
||||
)
|
||||
button_util_training_dir_output = gr.Button(
|
||||
"📂", 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")
|
||||
button_prepare_training_data.click(
|
||||
dreambooth_folder_preparation,
|
||||
inputs=[
|
||||
util_training_images_dir_input,
|
||||
util_training_images_repeat_input,
|
||||
util_instance_prompt_input,
|
||||
util_regularization_images_dir_input,
|
||||
util_regularization_images_repeat_input,
|
||||
util_class_prompt_input,
|
||||
util_training_dir_output,
|
||||
],
|
||||
)
|
||||
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
|
||||
])
|
||||
Loading…
Reference in New Issue
Block a user