From a549df730f7d3e5ecb5f95a225b200a654829ca0 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 16 Sep 2024 10:10:35 +0300
Subject: [PATCH] better caption check

---
 library/train_util.py | 2 ++
 nodes.py              | 5 +++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index a960522..4af4f5b 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1784,6 +1784,8 @@ class DreamBoothDataset(BaseDataset):
 
             self.set_tag_frequency(os.path.basename(subset.image_dir), captions)  # タグ頻度を記録
 
+            logger.info(f"Found captions for {len(captions) - len(missing_captions)} images.")
+            
             if missing_captions:
                 number_of_missing_captions = len(missing_captions)
                 number_of_missing_captions_to_show = 5
diff --git a/nodes.py b/nodes.py
index 640ef2e..0e3cfcf 100644
--- a/nodes.py
+++ b/nodes.py
@@ -88,6 +88,7 @@ class TrainDatasetGeneralConfig:
             },
             "optional": {
                 "reset_on_queue": ("BOOLEAN",{"default": False, "tooltip": "Force refresh of everything for cleaner queueing"}),
+                "caption_extension": ("STRING",{"default": ".txt", "tooltip": "extension for caption files"}),
             }
         }
 
@@ -96,12 +97,12 @@ class TrainDatasetGeneralConfig:
     FUNCTION = "create_config"
     CATEGORY = "FluxTrainer"
 
-    def create_config(self, shuffle_caption, caption_dropout_rate, color_aug, flip_aug, alpha_mask, reset_on_queue=False):
+    def create_config(self, shuffle_caption, caption_dropout_rate, color_aug, flip_aug, alpha_mask, reset_on_queue=False, caption_extension=".txt"):
         
         dataset = {
            "general": {
                 "shuffle_caption": shuffle_caption,
-                "caption_extension": ".txt",
+                "caption_extension": caption_extension,
                 "keep_tokens_separator": "|||",
                 "caption_dropout_rate": caption_dropout_rate,
                 "color_aug": color_aug,