Update De_DiffusionV2_Image.py

De_DiffusionV2_Image.py

@@ -145,8 +145,15 @@ class SDModel(PreTrainedModel):
         un_context_embeddings = self.text_encoder(un_token).last_hidden_state
         un_context_embeddings = un_context_embeddings.expand(batch_size, -1, -1)
         if self.training_args.use_text_encoder:
+            # Create attention mask for conditional_context
+            context_attention_mask = torch.ones(
+                (batch_size, conditional_context.shape[1]),
+                dtype=torch.long,
+                device=self._device
+            )
             context_embeddings = self.text_encoder(
-                inputs_embeds=conditional_context.to(self._dtype)
+                inputs_embeds=conditional_context.to(self._dtype),
+                attention_mask=context_attention_mask
             ).last_hidden_state # 1, 77 , 1024
 
         latent_shape = (batch_size, 4, self.training_args.image_size // 8, self.training_args.image_size // 8)
@@ -250,7 +257,17 @@ class SDModel(PreTrainedModel):
         conditional_context = self.language_proj(conditional_context) # [b, 159, 1024]
 
         if self.training_args.use_text_encoder:
-            text_encoder_output = self.text_encoder(input_ids=None, inputs_embeds=conditional_context.to(self._dtype))
+            # Create attention mask for conditional_context
+            context_attention_mask = torch.ones(
+                (self.batch_size, conditional_context.shape[1]),
+                dtype=torch.long,
+                device=self._device
+            )
+            text_encoder_output = self.text_encoder(
+                input_ids=None, 
+                inputs_embeds=conditional_context.to(self._dtype),
+                attention_mask=context_attention_mask
+            )
             pred_noise = self._unet_pred_noise(x_start=latent, t=timesteps, noise=noise, context=text_encoder_output.last_hidden_state.to(self._dtype)).to(self._dtype)
         else:
             pred_noise = self._unet_pred_noise(x_start=latent, t=timesteps, noise=noise, context=conditional_context.to(self._dtype)).to(self._dtype)