Merge pull request #2415 from bmaltais/2402-always_first_tags-is-not-…

…passwed-when-using-wd14-captioning

Fix issue with pre and post fix caption in subfolders

README.md

@@ -457,7 +457,8 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG
 
 ### 2024/04/28 (v24.0.9)
 
-- Update tmp config file to carry date and time info as part of the name. This will allow for easy batching of multiple training commands for users that desire to do so.
+- Updated the temporary configuration file to include date and time information in the file name. This will allow for easier batching of multiple training commands, particularly useful for users who want to automate their training sessions.
+- Fixed an issue with wd14 captioning where the captioning process was not functioning correctly when the recursive option was set to true. Prefixes and postfixes are now applied to all caption files in the folder.
 
 ### 2024/04/26 (v24.0.8)
 

kohya_gui/common_gui.py

@@ -739,6 +739,8 @@ def add_pre_postfix(
     prefix: str = "",
     postfix: str = "",
     caption_file_ext: str = ".caption",
+    recursive: bool = False,
+
 ) -> None:
     """
     Add prefix and/or postfix to the content of caption files within a folder.
@@ -762,10 +764,18 @@ def add_pre_postfix(
     # Define the image file extensions to filter
     image_extensions = (".jpg", ".jpeg", ".png", ".webp")
 
-    # List all image files in the folder
-    image_files = [
-        f for f in os.listdir(folder) if f.lower().endswith(image_extensions)
-    ]
+    # If recursive is true, list all image files in the folder and its subfolders
+    if recursive:
+        image_files = []
+        for root, dirs, files in os.walk(folder):
+            for file in files:
+                if file.lower().endswith(image_extensions):
+                    image_files.append(os.path.join(root, file))
+    else:
+        # List all image files in the folder
+        image_files = [
+            f for f in os.listdir(folder) if f.lower().endswith(image_extensions)
+        ]
 
     # Iterate over the list of image files
     for image_file in image_files:

kohya_gui/wd14_caption_gui.py

@@ -63,7 +63,7 @@ def caption_images(
     ]
 
     # Uncomment and modify if needed
-    # if always_first_tags:
+    # if always_first_tags != "":
     #     run_cmd.append('--always_first_tags')
     #     run_cmd.append(always_first_tags)
 
@@ -138,6 +138,7 @@ def caption_images(
         folder=train_data_dir,
         caption_file_ext=caption_extension,
         prefix=always_first_tags,
+        recursive=recursive,
     )
 
     log.info("...captioning done")
@@ -267,7 +268,7 @@ def list_train_dirs(path):
         with gr.Row():
             always_first_tags = gr.Textbox(
                 label="Prefix to add to WD14 caption",
-                info="comma-separated list of tags to always put at the beginning, e.g. 1girl, 1boy, ",
+                info="comma-separated list of tags to always put at the beginning, e.g.: 1girl, 1boy, ",
                 placeholder="(Optional)",
                 interactive=True,
                 value=config.get("wd14_caption.always_first_tags", ""),