v17.2 (#8)

* Update youtube video

* Dataset balancing

* Fix typo

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.

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')
 

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.'
         )

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 <num>_<text> 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],
+        )