Fix bug in stable diffusion when mask_pad_tokens is false (mosaicml#147)

coryMosaicML · Jun 4, 2024 · 8962008 · 8962008
1 parent b4d7380
commit 8962008
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/diffusion/models/stable_diffusion.py b/diffusion/models/stable_diffusion.py
@@ -484,7 +484,10 @@ def _prepare_text_embeddings(self, prompt, tokenized_prompts, tokenized_pad_mask
                                                truncation=True,
                                                return_tensors='pt')
                 tokenized_prompts = tokenized_out['input_ids']
-                tokenized_pad_mask = tokenized_out['attention_mask']
+                if self.mask_pad_tokens:
+                    tokenized_pad_mask = tokenized_out['attention_mask']
+                else:
+                    tokenized_pad_mask = None
             if tokenized_pad_mask is not None:
                 tokenized_pad_mask = tokenized_pad_mask.to(device)
             text_encoder_out = self.text_encoder(tokenized_prompts.to(device), attention_mask=tokenized_pad_mask)