modality mask

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1177,8 +1177,8 @@ def scan_fn(carry, block_and_mask):
             temb_prompt_audio=temb_prompt_audio,
             video_rotary_emb=video_rotary_emb,
             audio_rotary_emb=audio_rotary_emb,
-            ca_video_rotary_emb=ca_video_rotary_emb,
-            ca_audio_rotary_emb=ca_audio_rotary_emb,
+            ca_video_rotary_emb=video_cross_attn_rotary_emb,
+            ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
             a2v_cross_attention_mask=a2v_cross_attention_mask,
             v2a_cross_attention_mask=v2a_cross_attention_mask,
             attention_mask=mask,