refactor: deduplicate LTX2 guidance paths

Author: Perseus14

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -55,7 +55,6 @@ class LTX2BlockContext:
   a2v_cross_attention_mask: Optional[jax.Array] = None
   v2a_cross_attention_mask: Optional[jax.Array] = None
   perturbation_mask: Optional[jax.Array] = None
-  cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None
 
 
 class LTX2AdaLayerNormSingle(nnx.Module):
@@ -470,6 +469,7 @@ def named_scope(self, name: str):
   def __call__(
       self,
       ctx: "LTX2BlockContext",
+      cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None,
   ) -> Tuple[jax.Array, jax.Array]:
     """
     Forward pass of the LTX2 video/audio transformer block.
@@ -512,7 +512,6 @@ def __call__(
     a2v_cross_attention_mask = ctx.a2v_cross_attention_mask
     v2a_cross_attention_mask = ctx.v2a_cross_attention_mask
     perturbation_mask = ctx.perturbation_mask
-    cached_kv = ctx.cached_kv
     batch_size = hidden_states.shape[0]
 
     axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
@@ -1466,20 +1465,17 @@ def __call__(
           a2v_cross_attention_mask=a2v_cross_attention_mask,
           v2a_cross_attention_mask=v2a_cross_attention_mask,
           modality_mask=modality_mask,
-          cached_kv=cached_kv,
       )
 
       def apply_block(block, context: LTX2BlockContext, mask, block_cached_kv=None):
         orig_perturbation_mask = context.perturbation_mask
-        orig_cached_kv = context.cached_kv
-        context = context.replace(perturbation_mask=mask, cached_kv=block_cached_kv)
+        context = context.replace(perturbation_mask=mask)
         with self.named_scope("Transformer Layer"):
-          hidden_states_out, audio_hidden_states_out = block(context)
+          hidden_states_out, audio_hidden_states_out = block(context, cached_kv=block_cached_kv)
         context = context.replace(
             hidden_states=hidden_states_out.astype(context.hidden_states.dtype),
             audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype),
             perturbation_mask=orig_perturbation_mask,
-            cached_kv=orig_cached_kv,
         )
         return context