From 94529293008684d8c90ebd6255c04052d2b71d52 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 4 Oct 2024 10:55:20 -0600 Subject: [PATCH] Apply a mask to the embeds for SD if using T5 encoder --- toolkit/stable_diffusion_model.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 541822b..8fb9eef 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -2056,11 +2056,16 @@ class StableDiffusion: self.text_encoder, prompt, truncate=not long_prompts, - max_length=77, # todo set this higher when not transfer learning + max_length=256, dropout_prob=dropout_prob ) + + # just mask the attention mask + prompt_attention_mask = attention_mask.unsqueeze(-1).expand(embeds.shape) + embeds = embeds * prompt_attention_mask.to(dtype=embeds.dtype, device=embeds.device) return PromptEmbeds( embeds, + # do we want attn mask here? # attention_mask=attention_mask, )