From e1d66e47f475a8642c6e4654e1a51832a0db974b Mon Sep 17 00:00:00 2001 From: bmaltais Date: Sat, 17 Dec 2022 16:22:34 -0500 Subject: [PATCH] v17.2 (#8) * Update youtube video * Dataset balancing * Fix typo --- README.md | 9 ++- dreambooth_gui.py | 6 +- dreambooth_gui/caption_gui.py | 6 +- dreambooth_gui/dataset_balancing.py | 116 ++++++++++++++++++++++++++++ 4 files changed, 129 insertions(+), 8 deletions(-) create mode 100644 dreambooth_gui/dataset_balancing.py diff --git a/README.md b/README.md index dbcf060..1be4e21 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ python .\dreambooth_gui.py You can find a screen cast on how to use the GUI at the following location: -https://youtu.be/RlvqEKj03WI +[![Video](https://img.youtube.com/vi/RlvqEKj03WI/maxresdefault.jpg)](https://www.youtube.com/watch?v=RlvqEKj03WI) ## Folders configuration @@ -124,8 +124,13 @@ my_asd_dog_dreambooth Drop by the discord server for support: https://discord.com/channels/1041518562487058594/1041518563242020906 -## Change history +## Contributors +- Lord of the universe - cacoe (twitter: @cac0e) + +## Change history +* 12/?? (v17.2) update: + - Adding new dataset balancing utility. * 12/17 (v17.1) update: - Adding GUI for kohya_ss called dreambooth_gui.py - removing support for `--finetuning` as there is now a dedicated python repo for that. `--fine-tuning` is still there behind the scene until kohya_ss remove it in a future code release. diff --git a/dreambooth_gui.py b/dreambooth_gui.py index b148c3a..b5dd057 100644 --- a/dreambooth_gui.py +++ b/dreambooth_gui.py @@ -10,10 +10,9 @@ 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.dataset_balancing import gradio_dataset_balancing_tab from dreambooth_gui.common_gui import ( get_folder_path, remove_doublequote, @@ -729,6 +728,7 @@ with interface: ) # Captionning tab gradio_caption_gui_tab() + gradio_dataset_balancing_tab() button_run = gr.Button('Train model') diff --git a/dreambooth_gui/caption_gui.py b/dreambooth_gui/caption_gui.py index d3c36da..de3808d 100644 --- a/dreambooth_gui/caption_gui.py +++ b/dreambooth_gui/caption_gui.py @@ -18,7 +18,7 @@ def caption_images( return print( - f'Captionning files in {images_dir_input} with {caption_text_input}...' + f'Captioning files in {images_dir_input} with {caption_text_input}...' ) run_cmd = f'python "tools/caption.py"' run_cmd += f' --caption_text="{caption_text_input}"' @@ -33,7 +33,7 @@ def caption_images( # Run the command subprocess.run(run_cmd) - print('...captionning done') + print('...captioning done') ### @@ -42,7 +42,7 @@ def caption_images( def gradio_caption_gui_tab(): - with gr.Tab('Captionning'): + with gr.Tab('Captioning'): gr.Markdown( 'This utility will allow the creation of caption files for each images in a folder.' ) diff --git a/dreambooth_gui/dataset_balancing.py b/dreambooth_gui/dataset_balancing.py new file mode 100644 index 0000000..92d66c0 --- /dev/null +++ b/dreambooth_gui/dataset_balancing.py @@ -0,0 +1,116 @@ +import os +import re +import gradio as gr +from easygui import msgbox, boolbox +from .common_gui import get_folder_path + +# def select_folder(): +# # Open a file dialog to select a directory +# folder = filedialog.askdirectory() + +# # Update the GUI to display the selected folder +# selected_folder_label.config(text=folder) + + +def dataset_balancing(concept_repeats, folder, insecure): + + if not concept_repeats > 0: + # Display an error message if the total number of repeats is not a valid integer + msgbox('Please enter a valid integer for the total number of repeats.') + return + + concept_repeats = int(concept_repeats) + + # Check if folder exist + if folder == '' or not os.path.isdir(folder): + msgbox('Please enter a valid folder for balancing.') + return + + pattern = re.compile(r'^\d+_.+$') + + # Iterate over the subdirectories in the selected folder + for subdir in os.listdir(folder): + if pattern.match(subdir) or insecure: + # Calculate the number of repeats for the current subdirectory + # Get a list of all the files in the folder + files = os.listdir(os.path.join(folder, subdir)) + + # Filter the list to include only image files + image_files = [ + f + for f in files + if f.endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp')) + ] + + # Count the number of image files + images = len(image_files) + + # Check if the subdirectory name starts with a number inside braces, + # indicating that the repeats value should be multiplied + match = re.match(r'^\{(\d+\.?\d*)\}', subdir) + if match: + # Multiply the repeats value by the number inside the braces + repeats = max( + 1, round(concept_repeats / images * float(match.group(1))) + ) + subdir = subdir[match.end() :] + else: + repeats = max(1, round(concept_repeats / images)) + + # Check if the subdirectory name already has a number at the beginning + match = re.match(r'^\d+_', subdir) + if match: + # Replace the existing number with the new number + old_name = os.path.join(folder, subdir) + new_name = os.path.join( + folder, f'{repeats}_{subdir[match.end():]}' + ) + else: + # Add the new number at the beginning of the name + old_name = os.path.join(folder, subdir) + new_name = os.path.join(folder, f'{repeats}_{subdir}') + + os.rename(old_name, new_name) + else: + print(f"Skipping folder {subdir} because it does not match kohya_ss expected syntax...") + + msgbox('Dataset balancing completed...') + +def warning(insecure): + if insecure: + if boolbox(f'WARNING!!! You have asked to rename non kohya_ss _ folders...\n\nAre you sure you want to do that?', choices=("Yes, I like danger", "No, get me out of here")): + return True + else: + return False + +def gradio_dataset_balancing_tab(): + with gr.Tab('Dataset balancing'): + gr.Markdown('This utility will ensure that each concept folder in the dataset folder is used equally during the training process of the dreambooth machine learning model, regardless of the number of images in each folder. It will do this by renaming the concept folders to indicate the number of times they should be repeated during training.') + gr.Markdown('WARNING! The use of this utility on the wrong folder can lead to unexpected folder renaming!!!') + with gr.Row(): + select_dataset_folder_input = gr.Textbox(label="Dataset folder", + placeholder='Folder containing the concepts folders to balance...', + interactive=True, + ) + + select_dataset_folder_button = gr.Button( + '📂', elem_id='open_folder_small' + ) + select_dataset_folder_button.click( + get_folder_path, outputs=select_dataset_folder_input + ) + + total_repeats_number = gr.Number( + value=1000, + min=1, + interactive=True, + label='Training steps per concept per epoch', + ) + with gr.Accordion('Advanced options', open=False): + insecure = gr.Checkbox(value=False, label="DANGER!!! -- Insecure folder renaming -- DANGER!!!") + insecure.change(warning, inputs=insecure, outputs=insecure) + balance_button = gr.Button('Balance dataset') + balance_button.click( + dataset_balancing, + inputs=[total_repeats_number, select_dataset_folder_input, insecure], + )