diff --git a/dreambooth_gui.py b/dreambooth_gui.py index e2b515a..a681bf1 100644 --- a/dreambooth_gui.py +++ b/dreambooth_gui.py @@ -1,6 +1,7 @@ # v1: initial release # v2: add open and save folder icons # v3: Add new Utilities tab for Dreambooth folder preparation +# v3.1: Adding captionning of images to utilities import gradio as gr import json @@ -9,6 +10,7 @@ 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 @@ -475,29 +477,38 @@ def dreambooth_folder_preparation( util_regularization_images_dir_input, util_regularization_images_repeat_input, util_class_prompt_input, - util_training_dir_input, + util_training_dir_output, ): # Check if the input variables are empty - if (not len(util_training_dir_input)): + 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_input directory if it doesn't exist - os.makedirs(util_training_dir_input, exist_ok=True) + # 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 (not len(util_instance_prompt_input) - or not util_training_images_repeat_input > 0): + if (util_training_images_dir_input == ""): print( - "Training images directory or repeats is missing... can't perform the required task..." + "Training images directory is missing... can't perform the required task..." ) return else: training_dir = os.path.join( - util_training_dir_input, + util_training_dir_output, f"img/{int(util_training_images_repeat_input)}_{util_instance_prompt_input} {util_class_prompt_input}", ) @@ -518,7 +529,7 @@ def dreambooth_folder_preparation( ) else: regularization_dir = os.path.join( - util_training_dir_input, + util_training_dir_output, f"reg/{int(util_regularization_images_repeat_input)}_{util_class_prompt_input}", ) @@ -535,7 +546,7 @@ def dreambooth_folder_preparation( regularization_dir) print( - f"Done creating kohya_ss training folder structure at {util_training_dir_input}..." + f"Done creating kohya_ss training folder structure at {util_training_dir_output}..." ) @@ -810,16 +821,16 @@ with interface: interactive=True, elem_id="number_input") with gr.Row(): - util_training_dir_input = gr.Textbox( + 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_input = gr.Button( + button_util_training_dir_output = gr.Button( "📂", elem_id="open_folder_small") - button_util_training_dir_input.click( - get_folder_path, outputs=util_training_dir_input) + 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, @@ -830,16 +841,17 @@ with interface: util_regularization_images_dir_input, util_regularization_images_repeat_input, util_class_prompt_input, - util_training_dir_input, + util_training_dir_output, ], ) button_copy_info_to_Directories_tab = gr.Button( "Copy info to Directories Tab") + caption_gui.gradio_caption_gui() button_run = gr.Button("Train model") button_copy_info_to_Directories_tab.click(copy_info_to_Directories_tab, - inputs=[util_training_dir_input], + inputs=[util_training_dir_output], outputs=[ train_data_dir_input, reg_data_dir_input, diff --git a/dreambooth_gui/__init__.py b/dreambooth_gui/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dreambooth_gui/caption_gui.py b/dreambooth_gui/caption_gui.py new file mode 100644 index 0000000..21372c3 --- /dev/null +++ b/dreambooth_gui/caption_gui.py @@ -0,0 +1,75 @@ +import gradio as gr +from easygui import diropenbox, msgbox +import subprocess + +def get_folder_path(): + folder_path = diropenbox("Select the directory to use") + + return folder_path + +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...") + return + + # Check for images_dir_input + if images_dir_input == "": + msgbox("Image folder is missing...") + return + + 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 != "": + 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") + +### +# Gradio UI +### + +def gradio_caption_gui(): + with gr.Tab("Captionning"): + gr.Markdown( + "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", + interactive=True, + ) + overwrite_input = gr.Checkbox( + label="Overwrite existing captions in folder", + interactive=True, + value=False + ) + caption_file_ext = gr.Textbox( + 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", + interactive=True, + ) + button_images_dir_input = gr.Button( + "📂", 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]) + \ No newline at end of file