removed redundance img attn mask

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -1070,6 +1070,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:
@@ -1129,9 +1130,13 @@ 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 (created in embeddings_flax.py)
+        # It contains the image mask: [1]*257 + [0]*127 for 257 real image tokens padded to 384
+        if encoder_attention_mask is not None:
+            encoder_attention_mask_img = encoder_attention_mask[:, :padded_img_len]
+        else:
+            # Fallback: no mask means treat all as valid
+            encoder_attention_mask_img = None
       else:
         # If no image_seq_len is specified, treat all as text
         encoder_hidden_states_img = None

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -373,6 +373,7 @@ def __call__(
       rotary_emb: jax.Array,
       deterministic: bool = True,
       rngs: nnx.Rngs = None,
+      encoder_attention_mask: Optional[jax.Array] = None,
   ):
     with self.conditional_named_scope("transformer_block"):
       shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = jnp.split(
@@ -409,6 +410,7 @@ def __call__(
               encoder_hidden_states=encoder_hidden_states,
               deterministic=deterministic,
               rngs=rngs,
+              encoder_attention_mask = encoder_attention_mask
           )
         with self.conditional_named_scope("cross_attn_residual"):
           hidden_states = hidden_states + attn_output
@@ -621,7 +623,7 @@ def __call__(
       def scan_fn(carry, block):
         hidden_states_carry, rngs_carry = carry
         hidden_states = block(
-            hidden_states_carry, encoder_hidden_states, timestep_proj, rotary_emb, deterministic, rngs_carry
+            hidden_states_carry, encoder_hidden_states, timestep_proj, rotary_emb, deterministic, rngs_carry, encoder_attention_mask
         )
         new_carry = (hidden_states, rngs_carry)
         return new_carry, None