Upload matryoshka.py

unet/matryoshka.py

@@ -3059,6 +3059,7 @@ class MatryoshkaUNet2DConditionModel(
         added_cond_kwargs["masked_cross_attention"] = self.config.masked_cross_attention
         added_cond_kwargs["micro_conditioning_scale"] = self.config.micro_conditioning_scale
         added_cond_kwargs["from_nested"] = from_nested
+        added_cond_kwargs["conditioning_mask"] = encoder_attention_mask
 
         if not from_nested:
             encoder_hidden_states = self.process_encoder_hidden_states(
@@ -3507,6 +3508,7 @@ class NestedUNet2DConditionModel(MatryoshkaUNet2DConditionModel):
             added_cond_kwargs = added_cond_kwargs or {}
             added_cond_kwargs["masked_cross_attention"] = self.inner_unet.config.masked_cross_attention
             added_cond_kwargs["micro_conditioning_scale"] = self.config.micro_conditioning_scale
+            added_cond_kwargs["conditioning_mask"] = encoder_attention_mask
 
             if not self.config.nesting:
                 encoder_hidden_states = self.inner_unet.process_encoder_hidden_states(
@@ -3529,6 +3531,7 @@ class NestedUNet2DConditionModel(MatryoshkaUNet2DConditionModel):
             added_cond_kwargs = added_cond_kwargs or {}
             added_cond_kwargs["masked_cross_attention"] = self.inner_unet.inner_unet.config.masked_cross_attention
             added_cond_kwargs["micro_conditioning_scale"] = self.config.micro_conditioning_scale
+            added_cond_kwargs["conditioning_mask"] = encoder_attention_mask
 
             encoder_hidden_states = self.inner_unet.inner_unet.process_encoder_hidden_states(
                 encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
@@ -3603,7 +3606,7 @@ class NestedUNet2DConditionModel(MatryoshkaUNet2DConditionModel):
                     encoder_hidden_states=encoder_hidden_states[:bh],
                     attention_mask=attention_mask,
                     cross_attention_kwargs=cross_attention_kwargs,
-                    encoder_attention_mask=cond_mask[:bh] if cond_mask is not None else cond_mask,
+                    encoder_attention_mask=cond_mask_inner_unet[:bh] if cond_mask_inner_unet is not None else cond_mask_inner_unet,
                     **additional_residuals,
                 )
             else:
@@ -4025,7 +4028,7 @@ class MatryoshkaPipeline(
                 # Retrieve the original scale by scaling back the LoRA layers
                 unscale_lora_layers(self.text_encoder, lora_scale)
 
-        return prompt_embeds, negative_prompt_embeds
+        return prompt_embeds, negative_prompt_embeds, attention_mask
 
     def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
         dtype = next(self.image_encoder.parameters()).dtype
@@ -4458,7 +4461,7 @@ class MatryoshkaPipeline(
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
 
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+        prompt_embeds, negative_prompt_embeds, encoder_attention_mask = self.encode_prompt(
             prompt,
             device,
             num_images_per_prompt,
@@ -4548,6 +4551,7 @@ class MatryoshkaPipeline(
                     timestep_cond=timestep_cond,
                     cross_attention_kwargs=self.cross_attention_kwargs,
                     added_cond_kwargs=added_cond_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
                     return_dict=False,
                 )[0]