debug

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1193,8 +1193,8 @@ def scan_fn(carry, block_mask_and_id):
             audio_rotary_emb=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=None,
-            v2a_cross_attention_mask=None,
+            a2v_cross_attention_mask=encoder_attention_mask,
+            v2a_cross_attention_mask=audio_encoder_attention_mask,
             perturbation_mask=mask,
             modality_mask=modality_mask,
         )

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -833,9 +833,6 @@ def _print_stats(name, tensor):
         print(
             f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 4)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
         )
-      _print_stats("text_encoder_output_layer_0", prompt_embeds_list[0])
-      _print_stats("text_encoder_output_layer_last", prompt_embeds_list[-1])
-
       prompt_embeds = prompt_embeds_list
       del text_encoder_hidden_states  # Free PyTorch tensor memory
 
@@ -1949,8 +1946,7 @@ def scan_body(carry, inputs):
       )
 
       is_first_step = (t == timesteps_jax[0])
-      jax.lax.cond(is_first_step, lambda: print_stats_jit("noise_pred", noise_pred), lambda: None)
-      jax.lax.cond(is_first_step, lambda: print_stats_jit("noise_pred_audio", noise_pred_audio), lambda: None)
+
 
       # Extract latents_step based on stacking strategy
       if do_cfg and do_stg:
@@ -1984,6 +1980,7 @@ def scan_body(carry, inputs):
           x0_combined = rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale)
 
         noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t)
+        jax.lax.cond(is_first_step, lambda: print_stats_jit("noise_pred_video_after_guidance", noise_pred), lambda: None)
 
         # Audio guidance
         noise_pred_audio_uncond, noise_pred_audio_text, noise_pred_audio_perturb, noise_pred_audio_isolated = jnp.split(noise_pred_audio, 4, axis=0)
@@ -2003,6 +2000,7 @@ def scan_body(carry, inputs):
           x0_audio_combined = rescale_noise_cfg(x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale)
 
         noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t)
+        jax.lax.cond(is_first_step, lambda: print_stats_jit("noise_pred_audio_after_guidance", noise_pred_audio), lambda: None)
 
       # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic)
       elif do_cfg: