text attn fix

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -1110,7 +1110,10 @@ 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, attention_mask=attention_mask)
+        if is_self_attention:
+          attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
+        else:
+          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