diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 9bf569ff..c514ec33 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -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() diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 1a68720d..2a58bba8 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -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