Honor keep_tokens_separator / caption_separator when recording tag frequency metadata

library/dataset.py

@@ -478,11 +478,14 @@ def set_current_step(self, step):
     def set_max_train_steps(self, max_train_steps):
         self.max_train_steps = max_train_steps
 
-    def set_tag_frequency(self, dir_name, captions):
+    def set_tag_frequency(self, dir_name, captions, caption_separator=",", keep_tokens_separator=None):
         frequency_for_dir = self.tag_frequency.get(dir_name, {})
         self.tag_frequency[dir_name] = frequency_for_dir
         for caption in captions:
-            for tag in caption.split(","):
+            # split tags the same way as training does: keep_tokens_separator also separates tags
+            if keep_tokens_separator:
+                caption = caption.replace(keep_tokens_separator, caption_separator)
+            for tag in caption.split(caption_separator):
                 tag = tag.strip()
                 if tag:
                     tag = tag.lower()

library/dreambooth_dataset.py

@@ -242,7 +242,9 @@ def load_dreambooth_dir(subset: DreamBoothSubset):
                         else:
                             captions.append(cap_for_img)
 
-            self.set_tag_frequency(os.path.basename(subset.image_dir), captions)  # タグ頻度を記録
+            self.set_tag_frequency(
+                os.path.basename(subset.image_dir), captions, subset.caption_separator, subset.keep_tokens_separator
+            )  # タグ頻度を記録
 
             if missing_captions:
                 number_of_missing_captions = len(missing_captions)

library/finetuning_dataset.py

@@ -246,6 +246,8 @@ def __init__(
             self.num_train_images += len(metadata) * subset.num_repeats
 
             # TODO do not record tag freq when no tag
-            self.set_tag_frequency(os.path.basename(subset.metadata_file), tags_list)
+            self.set_tag_frequency(
+                os.path.basename(subset.metadata_file), tags_list, subset.caption_separator, subset.keep_tokens_separator
+            )
             subset.img_count = len(metadata)
             self.subsets.append(subset)

tests/library/test_tag_frequency.py

@@ -0,0 +1,23 @@
+from types import SimpleNamespace
+
+from library.dataset import BaseDataset
+
+
+def _frequency(captions, **kwargs):
+    dataset = SimpleNamespace(tag_frequency={})
+    BaseDataset.set_tag_frequency(dataset, "dir", captions, **kwargs)
+    return dataset.tag_frequency["dir"]
+
+
+def test_default_comma_separator_is_unchanged():
+    assert _frequency(["1girl, Solo", "1girl"]) == {"1girl": 2, "solo": 1}
+
+
+def test_keep_tokens_separator_splits_tags():
+    freq = _frequency(["1boy ||| standing, red", "1boy, blue ||| smile"], keep_tokens_separator="|||")
+    assert freq == {"1boy": 2, "standing": 1, "red": 1, "blue": 1, "smile": 1}
+
+
+def test_custom_caption_separator():
+    freq = _frequency(["1girl; solo ### smile"], caption_separator=";", keep_tokens_separator="###")
+    assert freq == {"1girl": 1, "solo": 1, "smile": 1}