From f421542df40aead2c0050f05cab1bff7ee7c07b5 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 11 Aug 2026 06:59:20 -0600 Subject: [PATCH] When dropping out, caching, and doing DOP, make sure we select a cahced trigger word when dropped out so DOP matches the drop out embeddings. --- toolkit/advanced_prompt_embeds.py | 18 +++++++++++++- toolkit/dataloader_mixins.py | 39 ++++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/toolkit/advanced_prompt_embeds.py b/toolkit/advanced_prompt_embeds.py index b86adeb..5dfc333 100644 --- a/toolkit/advanced_prompt_embeds.py +++ b/toolkit/advanced_prompt_embeds.py @@ -1,5 +1,6 @@ import os import torch +from safetensors import safe_open from safetensors.torch import load_file, save_file @@ -137,6 +138,8 @@ class AdvancedPromptEmbeds: def save(self, path): data = {} metadata = {"class_name": self.__class__.__name__} + if self._frozen_dtype_keys: + metadata["frozen_dtype_keys"] = ",".join(self._frozen_dtype_keys) for key, value in self._store.items(): if len(value) != 1: raise ValueError( @@ -150,6 +153,8 @@ class AdvancedPromptEmbeds: def load(cls, path=None): if path is not None: loaded = load_file(path) + with safe_open(path, framework="pt") as f: + metadata = f.metadata() else: raise ValueError("Must provide a path") @@ -157,7 +162,18 @@ class AdvancedPromptEmbeds: for key in loaded.keys(): data[key] = loaded[key] - return cls(**data) + pe = cls(**data) + frozen = [] + if metadata is not None and metadata.get("frozen_dtype_keys"): + frozen = metadata["frozen_dtype_keys"].split(",") + # files cached before frozen keys were saved have no metadata for them. + # Integer/bool tensors (token tags, etc) must never be dtype cast, so + # freeze them regardless. + for key, value in pe._store.items(): + if key not in frozen and not value[0].is_floating_point(): + frozen.append(key) + pe.frozen_dtype_keys = frozen + return pe @classmethod def concat_prompt_embeds( diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 6bb6310..1a68720 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -2092,7 +2092,10 @@ class TextEmbeddingFileItemDTOMixin: self._dop_text_embedding_path: Union[str, None] = None # blank caption embeds used for caption dropout when caching text embeddings self._blank_text_embedding_path: Union[str, None] = None + # DOP embeds for dropout steps (dropout caption with trigger replaced by class) + self._dop_blank_text_embedding_path: Union[str, None] = None self._loaded_text_embedding_path: Union[str, None] = None + self._caption_was_dropped = False self.is_text_embedding_cached = False self.text_embedding_load_device = 'cpu' self.text_embedding_version = 1 @@ -2161,6 +2164,28 @@ class TextEmbeddingFileItemDTOMixin: return inject_trigger_into_prompt('', trigger=self.trigger_word, add_if_not_present=True) return '' + def get_dop_dropout_caption(self: 'FileItemDTO'): + # live encoding replaces the trigger word with the preservation class on the + # dropped caption (class only), so the cached DOP dropout caption must match + dropout_caption = self.get_dropout_caption() + if self.trigger_word is not None: + return dropout_caption.replace( + self.trigger_word, self.dataset_config.diff_output_preservation_class + ) + return dropout_caption + + def get_dop_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False): + if self._dop_blank_text_embedding_path is not None and not recalculate: + return self._dop_blank_text_embedding_path + else: + # if the DOP dropout caption matches the dropout caption, this hashes to + # the same path as the blank embedding and the cache file is shared + self._dop_blank_text_embedding_path = self._build_text_embedding_path( + caption_override=self.get_dop_dropout_caption() + ) + + return self._dop_blank_text_embedding_path + def get_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False): if self._blank_text_embedding_path is not None and not recalculate: return self._blank_text_embedding_path @@ -2185,17 +2210,24 @@ class TextEmbeddingFileItemDTOMixin: return if self.prompt_embeds is None: te_path = self.get_text_embedding_path() + self._caption_was_dropped = False if self.dataset_config.caption_dropout_rate > 0: # get a random float form 0 to 1 rand = random.random() if rand < self.dataset_config.caption_dropout_rate: # drop the caption by using the cached blank embedding te_path = self.get_blank_text_embedding_path() + self._caption_was_dropped = True # load it from disk self.prompt_embeds = PromptEmbeds.load(te_path) self._loaded_text_embedding_path = te_path if self.dataset_config.diff_output_preservation and self.dop_prompt_embeds is None: - dop_path = self.get_dop_text_embedding_path() + if self._caption_was_dropped: + # match live encoding, which builds the DOP caption from the + # dropped caption (trigger word replaced with the class) + dop_path = self.get_dop_blank_text_embedding_path() + else: + dop_path = self.get_dop_text_embedding_path() if dop_path == self._loaded_text_embedding_path: # no trigger word in caption, same embedding self.dop_prompt_embeds = self.prompt_embeds @@ -2234,6 +2266,11 @@ class TextEmbeddingCachingMixin: if blank_path != text_embedding_path: # cache the dropout caption embedding (blank, or trigger word only) encode_targets.append((blank_path, file_item.get_dropout_caption())) + if self.dataset_config.diff_output_preservation: + # cache the DOP version of the dropout caption (class only) + dop_blank_path = file_item.get_dop_blank_text_embedding_path(recalculate=True) + if dop_blank_path not in [t[0] for t in encode_targets] + [text_embedding_path]: + encode_targets.append((dop_blank_path, file_item.get_dop_dropout_caption())) # only process if not saved to disk encode_targets = [t for t in encode_targets if not os.path.exists(t[0])] if len(encode_targets) > 0: