RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues 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.

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
extensions_built_in/sd_trainer/SDTrainer.py
View File

@ -120,6 +120,24 @@ class SDTrainer(BaseSDTrainProcess):
def before_model_load(self):
pass
def get_blank_control_image(self):
# noise instead of a black image so the fallback does not read as a
# meaningful (solid black) reference
control_image = torch.rand((1, 3, 224, 224), device=self.sd.device_torch, dtype=self.sd.torch_dtype)
if self.sd.has_multiple_control_images:
control_image = [control_image]
return control_image
def encode_static_prompt(self, prompt, **kwargs):
# static embeds (blank/trigger/uncond/DOP class) are always plain text.
# Some models (edit models) cannot encode a prompt without control images
# and raise, only then fall back to a blank control image. Real errors
# surface on the fallback call.
try:
return self.sd.encode_prompt(prompt, **kwargs)
except Exception:
return self.sd.encode_prompt(prompt, control_images=self.get_blank_control_image(), **kwargs)
def cache_sample_prompts(self):
if self.train_config.disable_sampling:
@ -259,18 +277,9 @@ class SDTrainer(BaseSDTrainProcess):
# cache unconditional embeds (blank prompt)
with torch.no_grad():
kwargs = {}
if self.sd.encode_control_in_text_embeddings:
# just do a blank image for unconditionals
control_image = torch.zeros((1, 3, 224, 224), device=self.sd.device_torch, dtype=self.sd.torch_dtype)
if self.sd.has_multiple_control_images:
control_image = [control_image]
kwargs['control_images'] = control_image
self.unconditional_embeds = self.sd.encode_prompt(
self.unconditional_embeds = self.encode_static_prompt(
[self.train_config.unconditional_prompt],
long_prompts=self.do_long_prompts,
**kwargs
).to(
self.device_torch,
dtype=self.sd.torch_dtype
@ -324,18 +333,11 @@ class SDTrainer(BaseSDTrainProcess):
raise ValueError("Cannot unload text encoder if training text encoder")
# cache embeddings
self.sd.text_encoder_to(self.device_torch)
encode_kwargs = {}
if self.sd.encode_control_in_text_embeddings:
# just do a blank image for unconditionals
control_image = torch.zeros((1, 3, 224, 224), device=self.sd.device_torch, dtype=self.sd.torch_dtype)
if self.sd.has_multiple_control_images:
control_image = [control_image]
encode_kwargs['control_images'] = control_image
self.cached_blank_embeds = self.sd.encode_prompt("", **encode_kwargs)
self.cached_blank_embeds = self.encode_static_prompt("")
if self.trigger_word is not None:
self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word, **encode_kwargs)
self.cached_trigger_embeds = self.encode_static_prompt(self.trigger_word)
if self.train_config.diff_output_preservation:
self.cached_dop_class_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class)
self.cached_dop_class_embeds = self.encode_static_prompt(self.train_config.diff_output_preservation_class)
self.diff_output_preservation_embeds = self.cached_dop_class_embeds
self.cache_sample_prompts()

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