testing flash with audio_to_video attn

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -239,7 +239,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.attention_kernel,
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -353,9 +353,6 @@ def __call__(
     axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed"))
     audio_hidden_states = jax.lax.with_sharding_constraint(audio_hidden_states, axis_names_audio)
 
-    jax.debug.print("[LTX2] Video hidden states shape: {shape}", shape=hidden_states.shape)
-    jax.debug.print("[LTX2] Audio hidden states shape: {shape}", shape=audio_hidden_states.shape)
-
     if encoder_hidden_states is not None:
       encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, axis_names)
     if audio_encoder_hidden_states is not None: