text attn mask fix

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -1072,6 +1072,7 @@ def __call__(
       hidden_states: jax.Array,
       encoder_hidden_states: jax.Array = None,
       rotary_emb: Optional[jax.Array] = None,
+      encoder_attention_mask: Optional[jax.Array] = None,
       deterministic: bool = True,
       rngs: nnx.Rngs = None,
   ) -> jax.Array:
@@ -1111,7 +1112,7 @@ def __call__(
       value_proj = checkpoint_name(value_proj, "value_proj")
 
       with jax.named_scope("apply_attention"):
-        attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
+        attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj, attention_mask=encoder_attention_mask)
 
     else:
       # NEW PATH for I2V CROSS-ATTENTION
@@ -1131,9 +1132,14 @@ def __call__(
         encoder_hidden_states_img = encoder_hidden_states[:, :padded_img_len, :]
         encoder_hidden_states_text = encoder_hidden_states[:, padded_img_len:, :]
 
-        encoder_attention_mask_img = jnp.ones((encoder_hidden_states_img.shape[0], padded_img_len), dtype=jnp.int32)
-        if image_seq_len_actual < padded_img_len:
-             encoder_attention_mask_img = encoder_attention_mask_img.at[:, image_seq_len_actual:].set(0)
+        # Use the passed encoder_attention_mask, which already contains both image and text masks
+        if encoder_attention_mask is not None:
+            encoder_attention_mask_img = encoder_attention_mask[:, :padded_img_len]
+            encoder_attention_mask_text = encoder_attention_mask[:, padded_img_len:]
+        else:
+            # Fallback: if no mask passed, treat all as valid (shouldn't happen with our fix)
+            encoder_attention_mask_img = None
+            encoder_attention_mask_text = None
       else:
         # If no image_seq_len is specified, treat all as text
         encoder_hidden_states_img = None
@@ -1176,7 +1182,7 @@ def __call__(
 
         # Attention - tensors are (B, S, D)
         with self.conditional_named_scope("cross_attn_text_apply"):
-          attn_output_text = self.attention_op.apply_attention(query_proj_text, key_proj_text, value_proj_text)
+          attn_output_text = self.attention_op.apply_attention(query_proj_text, key_proj_text, value_proj_text, attention_mask=encoder_attention_mask_text)
         with self.conditional_named_scope("cross_attn_img_apply"):
           # Pass encoder_attention_mask_img for image cross-attention to mask padded tokens
           attn_output_img = self.attention_op.apply_attention(query_proj_img, key_proj_img, value_proj_img, attention_mask=encoder_attention_mask_img)
@@ -1189,7 +1195,7 @@ def __call__(
         value_proj_text = checkpoint_name(value_proj_text, "value_proj_text")
 
         with self.conditional_named_scope("cross_attn_text_apply"):
-          attn_output = self.attention_op.apply_attention(query_proj_text, key_proj_text, value_proj_text)
+          attn_output = self.attention_op.apply_attention(query_proj_text, key_proj_text, value_proj_text, attention_mask=encoder_attention_mask_text)
 
     attn_output = attn_output.astype(dtype=dtype)
     attn_output = checkpoint_name(attn_output, "attn_output")