Address PR comments: Add cross attention masks to context and remove unused rngs_carry from apply_block

Author: Perseus14

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1166,6 +1166,12 @@ def __call__(
     # Determine timestep for audio.
     audio_timestep = audio_timestep if audio_timestep is not None else timestep
 
+    a2v_cross_attention_mask = None
+    v2a_cross_attention_mask = None
+    if attention_kwargs is not None:
+      a2v_cross_attention_mask = attention_kwargs.get("a2v_cross_attention_mask", None)
+      v2a_cross_attention_mask = attention_kwargs.get("v2a_cross_attention_mask", None)
+
     if self.attention_kernel == "dot_product":
       if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
         encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0
@@ -1175,6 +1181,14 @@ def __call__(
         audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0
         audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1)
 
+      if a2v_cross_attention_mask is not None and a2v_cross_attention_mask.ndim == 2:
+        a2v_cross_attention_mask = (1 - a2v_cross_attention_mask.astype(self.dtype)) * -10000.0
+        a2v_cross_attention_mask = jnp.expand_dims(a2v_cross_attention_mask, axis=1)
+
+      if v2a_cross_attention_mask is not None and v2a_cross_attention_mask.ndim == 2:
+        v2a_cross_attention_mask = (1 - v2a_cross_attention_mask.astype(self.dtype)) * -10000.0
+        v2a_cross_attention_mask = jnp.expand_dims(v2a_cross_attention_mask, axis=1)
+
     batch_size = hidden_states.shape[0]
 
     # 1. Prepare RoPE positional embeddings
@@ -1291,10 +1305,12 @@ def __call__(
           ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
           encoder_attention_mask=encoder_attention_mask,
           audio_encoder_attention_mask=audio_encoder_attention_mask,
+          a2v_cross_attention_mask=a2v_cross_attention_mask,
+          v2a_cross_attention_mask=v2a_cross_attention_mask,
           modality_mask=modality_mask,
       )
 
-      def apply_block(block, context: LTX2BlockContext, mask, rngs_carry):
+      def apply_block(block, context: LTX2BlockContext, mask):
         orig_perturbation_mask = context.perturbation_mask
         context = dataclasses.replace(context, perturbation_mask=mask)
         with jax.named_scope("Transformer Layer"):
@@ -1305,13 +1321,13 @@ def apply_block(block, context: LTX2BlockContext, mask, rngs_carry):
             audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype),
             perturbation_mask=orig_perturbation_mask,
         )
-        return context, rngs_carry
+        return context
 
       if perturbation_mask is None:
 
         def scan_fn_ltx2(carry, block):
           context, rngs_carry = carry
-          context, rngs_carry = apply_block(block, context, None, rngs_carry)
+          context = apply_block(block, context, None)
           return (context, rngs_carry), None
 
         if self.scan_layers:
@@ -1333,7 +1349,7 @@ def scan_fn_ltx2(carry, block):
         else:
           current_context = base_context
           for block in self.transformer_blocks:
-            current_context, _ = apply_block(block, current_context, None, None)
+            current_context = apply_block(block, current_context, None)
           hidden_states = current_context.hidden_states
           audio_hidden_states = current_context.audio_hidden_states
       else:
@@ -1346,7 +1362,7 @@ def scan_fn_ltx2(carry, block):
         def scan_fn_ltx23(carry, block_and_mask):
           block, mask = block_and_mask
           context, rngs_carry = carry
-          context, rngs_carry = apply_block(block, context, mask, rngs_carry)
+          context = apply_block(block, context, mask)
           return (context, rngs_carry), None
 
         if self.scan_layers:
@@ -1369,7 +1385,7 @@ def scan_fn_ltx23(carry, block_and_mask):
           current_context = base_context
           for i, block in enumerate(self.transformer_blocks):
             mask = perturbation_mask_per_layer[i] if perturbation_mask_per_layer is not None else None
-            current_context, _ = apply_block(block, current_context, mask, None)
+            current_context = apply_block(block, current_context, mask)
           hidden_states = current_context.hidden_states
           audio_hidden_states = current_context.audio_hidden_states