From 257da9b5869b44f1c55c94303719ba5e82f6f542 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 9 Aug 2026 22:13:49 -0600 Subject: [PATCH] Rework DOP so it works with caching text embeddings --- extensions_built_in/sd_trainer/SDTrainer.py | 32 ++++++- jobs/process/BaseSDTrainProcess.py | 7 ++ toolkit/config_modules.py | 9 +- toolkit/data_transfer_object/data_loader.py | 26 +++++- toolkit/dataloader_mixins.py | 94 ++++++++++++++++----- 5 files changed, 138 insertions(+), 30 deletions(-) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 9552e88f..ff4b3c31 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -81,12 +81,19 @@ class SDTrainer(BaseSDTrainProcess): self.cached_blank_embeds: Optional[PromptEmbeds] = None self.cached_trigger_embeds: Optional[PromptEmbeds] = None self.diff_output_preservation_embeds: Optional[PromptEmbeds] = None + # fallback class-only embeds for when the text encoder is unloaded and + # per item DOP embeds were not cached to disk + self.cached_dop_class_embeds: Optional[PromptEmbeds] = None self.dfe: Optional[DiffusionFeatureExtractor] = None self.unconditional_embeds = None if self.train_config.diff_output_preservation: - if self.trigger_word is None: + # datasets can have their own trigger words, the global one is copied to them if not set + has_dataset_trigger = any( + dataset.trigger_word is not None for dataset in self.dataset_configs + ) + if self.trigger_word is None and not has_dataset_trigger: raise ValueError("diff_output_preservation requires a trigger_word to be set") if self.network_config is None: raise ValueError("diff_output_preservation requires a network to be set") @@ -328,7 +335,8 @@ class SDTrainer(BaseSDTrainProcess): if self.trigger_word is not None: self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word, **encode_kwargs) if self.train_config.diff_output_preservation: - self.diff_output_preservation_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class) + self.cached_dop_class_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class) + self.diff_output_preservation_embeds = self.cached_dop_class_embeds self.cache_sample_prompts() @@ -1660,6 +1668,16 @@ class SDTrainer(BaseSDTrainProcess): [unconditional_embeds] * noisy_latents.shape[0] ) + if self.train_config.diff_output_preservation: + if batch.dop_prompt_embeds is not None: + # cached to disk with the trigger word replaced per dataset + self.diff_output_preservation_embeds = batch.dop_prompt_embeds.clone().detach().to( + self.device_torch, dtype=dtype + ) + else: + # no per item cache, fall back to the class only embeds + self.diff_output_preservation_embeds = self.cached_dop_class_embeds + if isinstance(self.adapter, CustomAdapter): self.adapter.is_unconditional_run = False @@ -1726,10 +1744,16 @@ class SDTrainer(BaseSDTrainProcess): self.adapter.is_unconditional_run = False if self.train_config.diff_output_preservation: - dop_prompts = [p.replace(self.trigger_word, self.train_config.diff_output_preservation_class) for p in conditioned_prompts] + # datasets can have their own trigger words, replace per item + def replace_trigger_with_class(prompt, file_item): + trigger = file_item.trigger_word if file_item.trigger_word is not None else self.trigger_word + if trigger is None: + return prompt + return prompt.replace(trigger, self.train_config.diff_output_preservation_class) + dop_prompts = [replace_trigger_with_class(p, fi) for p, fi in zip(conditioned_prompts, batch.file_items)] dop_prompts_2 = None if prompt_2 is not None: - dop_prompts_2 = [p.replace(self.trigger_word, self.train_config.diff_output_preservation_class) for p in prompt_2] + dop_prompts_2 = [replace_trigger_with_class(p, fi) for p, fi in zip(prompt_2, batch.file_items)] self.diff_output_preservation_embeds = self.sd.encode_prompt( dop_prompts, dop_prompts_2, dropout_prob=self.train_config.prompt_dropout_prob, diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index b6777eec..ef415dc5 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -153,6 +153,13 @@ class BaseSDTrainProcess(BaseTrainProcess): if self.train_config.cache_text_embeddings: for raw_dataset in raw_datasets: raw_dataset['cache_text_embeddings'] = True + + # pass diff output preservation to the datasets so the data loader can build + # and cache the DOP caption (dataset trigger word replaced with the class) + if self.train_config.diff_output_preservation and raw_datasets is not None: + for raw_dataset in raw_datasets: + raw_dataset['diff_output_preservation'] = True + raw_dataset['diff_output_preservation_class'] = self.train_config.diff_output_preservation_class if raw_datasets is not None and len(raw_datasets) > 0: for raw_dataset in raw_datasets: diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 20a1f28f..a3a51138 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -924,6 +924,10 @@ class DatasetConfig: self.default_caption: str = kwargs.get('default_caption', None) # trigger word for just this dataset self.trigger_word: str = kwargs.get('trigger_word', None) + # set automatically from the train config when diff output preservation is enabled. + # the dataset trigger word is replaced with the class in the caption for DOP embeddings + self.diff_output_preservation: bool = kwargs.get('diff_output_preservation', False) + self.diff_output_preservation_class: str = kwargs.get('diff_output_preservation_class', '') random_triggers = kwargs.get('random_triggers', []) # if they are a string, load them from a file if isinstance(random_triggers, str) and os.path.exists(random_triggers): @@ -1489,11 +1493,6 @@ def validate_configs( # see if any datasets are caching text embeddings is_caching_text_embeddings = any(dataset.cache_text_embeddings for dataset in dataset_configs) if is_caching_text_embeddings: - - # check if they are doing differential output preservation - if train_config.diff_output_preservation: - raise ValueError("Cannot use differential output preservation with caching text embeddings. Please set diff_output_preservation to False.") - # make sure they are all cached for dataset in dataset_configs: if not dataset.cache_text_embeddings: diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 0f104522..c1a46451 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -315,6 +315,8 @@ class DataLoaderBatchDTO: ) self.prompt_embeds: Union[PromptEmbeds, None] = None + # diff output preservation embeds (trigger word replaced with class) + self.dop_prompt_embeds: Union[PromptEmbeds, None] = None # if self.file_items[0].control_tensor is not None: # if any have a control tensor, we concatenate them if any([x.control_tensor is not None for x in self.file_items]): @@ -482,9 +484,31 @@ class DataLoaderBatchDTO: y.text_embeds = [y.text_embeds] prompt_embeds_list.append(y) padding_side = self.file_items[0].te_padding_side - + self.prompt_embeds = concat_prompt_embeds(prompt_embeds_list, padding_side=padding_side) + if any([x.dop_prompt_embeds is not None for x in self.file_items]): + # find one to use as a base + base_dop_prompt_embeds = None + for x in self.file_items: + if x.dop_prompt_embeds is not None: + base_dop_prompt_embeds = x.dop_prompt_embeds + break + dop_prompt_embeds_list = [] + for x in self.file_items: + if x.dop_prompt_embeds is None: + y = base_dop_prompt_embeds + else: + y = x.dop_prompt_embeds + if x.text_embedding_space_version == "zimage": + # z image needs to be a list if it is not already + if not isinstance(y.text_embeds, list): + y.text_embeds = [y.text_embeds] + dop_prompt_embeds_list.append(y) + padding_side = self.file_items[0].te_padding_side + + self.dop_prompt_embeds = concat_prompt_embeds(dop_prompt_embeds_list, padding_side=padding_side) + if any([x.audio_tensor is not None for x in self.file_items]): # find one to use as a base base_audio_tensor = None diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 398be98a..a926fbdf 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -321,6 +321,8 @@ class CaptionProcessingDTOMixin: self.raw_caption_short: str = None self.caption: str = None self.caption_short: str = None + # caption with the trigger word replaced by the diff output preservation class + self.caption_dop: str = None dataset_config: DatasetConfig = kwargs.get('dataset_config', None) self.extra_values: List[float] = dataset_config.extra_values @@ -367,6 +369,14 @@ class CaptionProcessingDTOMixin: self.caption = self.get_caption() if self.raw_caption_short is not None: self.caption_short = self.get_caption(short_caption=True) + if self.dataset_config.diff_output_preservation: + # replace this dataset's trigger word with the preservation class. + # do it on the final caption so token order matches the normal caption + self.caption_dop = self.caption + if self.trigger_word is not None: + self.caption_dop = self.caption.replace( + self.trigger_word, self.dataset_config.diff_output_preservation_class + ) def get_caption( self: 'FileItemDTO', @@ -2077,17 +2087,20 @@ class TextEmbeddingFileItemDTOMixin: super().__init__(*args, **kwargs) self.prompt_embeds: Union[PromptEmbeds, None] = None self._text_embedding_path: Union[str, None] = None + # diff output preservation embeds (caption with trigger word replaced by class) + self.dop_prompt_embeds: Union[PromptEmbeds, None] = None + self._dop_text_embedding_path: Union[str, None] = None self.is_text_embedding_cached = False self.text_embedding_load_device = 'cpu' self.text_embedding_version = 1 - def get_text_embedding_info_dict(self: 'FileItemDTO'): + def get_text_embedding_info_dict(self: 'FileItemDTO', caption_override=None): # make sure the caption is loaded here # TODO: we need a way to cache all the other features like trigger words, DOP, etc. For now, we need to throw an error if not compatible. if self.caption is None: self.load_caption() item = OrderedDict([ - ("caption", self.caption), + ("caption", self.caption if caption_override is None else caption_override), ("text_embedding_space_version", self.text_embedding_space_version), ("text_embedding_version", self.text_embedding_version), ]) @@ -2103,27 +2116,47 @@ class TextEmbeddingFileItemDTOMixin: item["first_frame_in_te"] = True return item + def _build_text_embedding_path(self: 'FileItemDTO', caption_override=None): + # we store text embeddings in a folder in same path as image called _text_embedding_cache + img_dir = os.path.dirname(self.path) + te_dir = os.path.join(img_dir, '_t_e_cache') + hash_dict = self.get_text_embedding_info_dict(caption_override=caption_override) + filename_no_ext = os.path.splitext(os.path.basename(self.path))[0] + # get base64 hash of md5 checksum of hash_dict + hash_input = json.dumps(hash_dict, sort_keys=True).encode('utf-8') + hash_str = base64.urlsafe_b64encode(hashlib.md5(hash_input).digest()).decode('ascii') + hash_str = hash_str.replace('=', '') + return os.path.join(te_dir, f'{filename_no_ext}_{hash_str}.safetensors') + def get_text_embedding_path(self: 'FileItemDTO', recalculate=False): if self._text_embedding_path is not None and not recalculate: return self._text_embedding_path else: - # we store text embeddings in a folder in same path as image called _text_embedding_cache - img_dir = os.path.dirname(self.path) - te_dir = os.path.join(img_dir, '_t_e_cache') - hash_dict = self.get_text_embedding_info_dict() - filename_no_ext = os.path.splitext(os.path.basename(self.path))[0] - # get base64 hash of md5 checksum of hash_dict - hash_input = json.dumps(hash_dict, sort_keys=True).encode('utf-8') - hash_str = base64.urlsafe_b64encode(hashlib.md5(hash_input).digest()).decode('ascii') - hash_str = hash_str.replace('=', '') - self._text_embedding_path = os.path.join(te_dir, f'{filename_no_ext}_{hash_str}.safetensors') + self._text_embedding_path = self._build_text_embedding_path() return self._text_embedding_path + def get_dop_text_embedding_path(self: 'FileItemDTO', recalculate=False): + if self._dop_text_embedding_path is not None and not recalculate: + return self._dop_text_embedding_path + else: + # make sure the caption is loaded so caption_dop is built + if self.caption is None: + self.load_caption() + # if the trigger word is not in the caption, this hashes to the same + # path as the normal embedding and the cache file is shared + self._dop_text_embedding_path = self._build_text_embedding_path( + caption_override=self.caption_dop + ) + + return self._dop_text_embedding_path + def cleanup_text_embedding(self): if self.prompt_embeds is not None: # we are caching on disk, don't save in memory self.prompt_embeds = None + if self.dop_prompt_embeds is not None: + self.dop_prompt_embeds = None def load_prompt_embedding(self, device=None): if not self.is_text_embedding_cached: @@ -2131,6 +2164,13 @@ class TextEmbeddingFileItemDTOMixin: if self.prompt_embeds is None: # load it from disk self.prompt_embeds = PromptEmbeds.load(self.get_text_embedding_path()) + if self.dataset_config.diff_output_preservation and self.dop_prompt_embeds is None: + dop_path = self.get_dop_text_embedding_path() + if dop_path == self.get_text_embedding_path(): + # no trigger word in caption, same embedding + self.dop_prompt_embeds = self.prompt_embeds + else: + self.dop_prompt_embeds = PromptEmbeds.load(dop_path) class TextEmbeddingCachingMixin: def __init__(self: 'AiToolkitDataset', **kwargs): @@ -2152,13 +2192,21 @@ class TextEmbeddingCachingMixin: file_item.latent_load_device = self.sd.device text_embedding_path = file_item.get_text_embedding_path(recalculate=True) + # (path, caption) pairs to encode for this item + encode_targets = [(text_embedding_path, file_item.caption)] + if self.dataset_config.diff_output_preservation: + dop_path = file_item.get_dop_text_embedding_path(recalculate=True) + if dop_path != text_embedding_path: + # trigger word was in the caption, cache the DOP version too + encode_targets.append((dop_path, file_item.caption_dop)) # only process if not saved to disk - if not os.path.exists(text_embedding_path): + encode_targets = [t for t in encode_targets if not os.path.exists(t[0])] + if len(encode_targets) > 0: # load if not loaded if not did_move: self.sd.set_device_state_preset('cache_text_encoder') did_move = True - + if file_item.encode_control_in_text_embeddings and file_item.control_path is not None: ctrl_img_list = [] control_path_list = file_item.control_path @@ -2185,7 +2233,10 @@ class TextEmbeddingCachingMixin: ctrl_img = ctrl_img_list[0] else: ctrl_img = ctrl_img_list - prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption, control_images=ctrl_img) + for path, caption in encode_targets: + prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img) + prompt_embeds.save(path) + del prompt_embeds elif ( getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) and self.dataset_config.do_i2v @@ -2205,13 +2256,16 @@ class TextEmbeddingCachingMixin: ) if self.sd.has_multiple_control_images: ctrl_img = [ctrl_img] - prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption, control_images=ctrl_img) + for path, caption in encode_targets: + prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img) + prompt_embeds.save(path) + del prompt_embeds file_item.tensor = None else: - prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption) - # save it - prompt_embeds.save(text_embedding_path) - del prompt_embeds + for path, caption in encode_targets: + prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption) + prompt_embeds.save(path) + del prompt_embeds file_item.is_text_embedding_cached = True i += 1 # restore device state