From fa141e7f5569b121f253ef32ccd524086c01a32d Mon Sep 17 00:00:00 2001 From: kv-248 Date: Thu, 24 Sep 2026 22:43:55 +0530 Subject: [PATCH] fix: honor keep_tokens_separator and caption_separator in tag frequency metadata (#2323) --- library/dataset.py | 7 +++++-- library/dreambooth_dataset.py | 4 +++- library/finetuning_dataset.py | 4 +++- tests/library/test_tag_frequency.py | 23 +++++++++++++++++++++++ 4 files changed, 34 insertions(+), 4 deletions(-) create mode 100644 tests/library/test_tag_frequency.py diff --git a/library/dataset.py b/library/dataset.py index c6cb296e51..ad58a69fe5 100644 --- a/library/dataset.py +++ b/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() diff --git a/library/dreambooth_dataset.py b/library/dreambooth_dataset.py index 2759a74300..51158f854d 100644 --- a/library/dreambooth_dataset.py +++ b/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) diff --git a/library/finetuning_dataset.py b/library/finetuning_dataset.py index 84fdfdb264..2611adaef5 100644 --- a/library/finetuning_dataset.py +++ b/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) diff --git a/tests/library/test_tag_frequency.py b/tests/library/test_tag_frequency.py new file mode 100644 index 0000000000..80281a775c --- /dev/null +++ b/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}