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,