github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Only add a blank control if the model has to have it, for unconditionals. Previously it always encoded a blank control on unconditional if the model could take one. Affects DOP and blank prompt preservations.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-11 08:25:23 -06:00
parent f421542df4
commit 21dc65972d
2 changed files with 59 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
toolkit/dataloader_mixins.py
View File

@@ -2100,7 +2100,7 @@ class TextEmbeddingFileItemDTOMixin:
self.text_embedding_load_device = 'cpu'
self.text_embedding_version = 1
def get_text_embedding_info_dict(self: 'FileItemDTO', caption_override=None):
def get_text_embedding_info_dict(self: 'FileItemDTO', caption_override=None, text_only=False):
# 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:
@@ -2110,6 +2110,10 @@ class TextEmbeddingFileItemDTOMixin:
("text_embedding_space_version", self.text_embedding_space_version),
("text_embedding_version", self.text_embedding_version),
])
# dropout embeds are encoded as plain text, keep control conditioning
# out of their cache key
if text_only:
return item
# if we have a control image, cache the path
if self.encode_control_in_text_embeddings and self.control_path is not None:
item["control_path"] = self.control_path
@@ -2122,11 +2126,11 @@ class TextEmbeddingFileItemDTOMixin:
item["first_frame_in_te"] = True
return item
def _build_text_embedding_path(self: 'FileItemDTO', caption_override=None):
def _build_text_embedding_path(self: 'FileItemDTO', caption_override=None, text_only=False):
# 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)
hash_dict = self.get_text_embedding_info_dict(caption_override=caption_override, text_only=text_only)
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')
@@ -2179,9 +2183,10 @@ class TextEmbeddingFileItemDTOMixin:
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
# the same path as the blank embedding and the cache file is shared.
# text_only: dropout embeds carry no control conditioning
self._dop_blank_text_embedding_path = self._build_text_embedding_path(
caption_override=self.get_dop_dropout_caption()
caption_override=self.get_dop_dropout_caption(), text_only=True
)
return self._dop_blank_text_embedding_path
@@ -2190,10 +2195,12 @@ class TextEmbeddingFileItemDTOMixin:
if self._blank_text_embedding_path is not None and not recalculate:
return self._blank_text_embedding_path
else:
# if the dropout caption matches the normal caption, this hashes to the
# same path as the normal embedding and the cache file is shared
# if the dropout caption matches the normal caption (and the item has no
# control conditioning), this hashes to the same path as the normal
# embedding and the cache file is shared.
# text_only: dropout embeds carry no control conditioning
self._blank_text_embedding_path = self._build_text_embedding_path(
caption_override=self.get_dropout_caption()
caption_override=self.get_dropout_caption(), text_only=True
)
return self._blank_text_embedding_path
@@ -2261,16 +2268,20 @@ class TextEmbeddingCachingMixin:
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))
# dropout embeds are encoded as plain text (no control images)
dropout_target_paths = set()
if self.dataset_config.caption_dropout_rate > 0:
blank_path = file_item.get_blank_text_embedding_path(recalculate=True)
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()))
dropout_target_paths.add(blank_path)
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()))
dropout_target_paths.add(dop_blank_path)
# 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:
@@ -2306,7 +2317,15 @@ class TextEmbeddingCachingMixin:
else:
ctrl_img = ctrl_img_list
for path, caption in encode_targets:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
if path in dropout_target_paths:
# dropout embeds are plain text. Only fall back to the
# control images if the model cannot encode without them
try:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
except Exception:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
else:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
prompt_embeds.save(path)
del prompt_embeds
elif (
@@ -2329,7 +2348,15 @@ class TextEmbeddingCachingMixin:
if self.sd.has_multiple_control_images:
ctrl_img = [ctrl_img]
for path, caption in encode_targets:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
if path in dropout_target_paths:
# dropout embeds are plain text. Only fall back to the
# control images if the model cannot encode without them
try:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
except Exception:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
else:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
prompt_embeds.save(path)
del prompt_embeds
file_item.tensor = None