From 90dad5471c85055374aee684b7da23b5ea5c66c8 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Fri, 16 Dec 2022 13:16:23 -0500 Subject: [PATCH] Refactor GUI code to be modular --- dreambooth_gui.py | 216 +------------------ dreambooth_gui/caption_gui.py | 10 +- dreambooth_gui/common_gui.py | 19 ++ dreambooth_gui/dreambooth_folder_creation.py | 179 +++++++++++++++ 4 files changed, 212 insertions(+), 212 deletions(-) create mode 100644 dreambooth_gui/common_gui.py create mode 100644 dreambooth_gui/dreambooth_folder_creation.py diff --git a/dreambooth_gui.py b/dreambooth_gui.py index a681bf1..176de8a 100644 --- a/dreambooth_gui.py +++ b/dreambooth_gui.py @@ -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=[ diff --git a/dreambooth_gui/caption_gui.py b/dreambooth_gui/caption_gui.py index 21372c3..2b2123d 100644 --- a/dreambooth_gui/caption_gui.py +++ b/dreambooth_gui/caption_gui.py @@ -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." diff --git a/dreambooth_gui/common_gui.py b/dreambooth_gui/common_gui.py new file mode 100644 index 0000000..61579b1 --- /dev/null +++ b/dreambooth_gui/common_gui.py @@ -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 \ No newline at end of file diff --git a/dreambooth_gui/dreambooth_folder_creation.py b/dreambooth_gui/dreambooth_folder_creation.py new file mode 100644 index 0000000..f2bdb0d --- /dev/null +++ b/dreambooth_gui/dreambooth_folder_creation.py @@ -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 + ]) \ No newline at end of file