debug

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1450,6 +1450,20 @@ def _print_stats(name, tensor):
         audio_embeds_sharded = jax.device_put(audio_embeds, spec)
       print(f"WEIGHT DEBUG: block 0 to_q kernel mean: {float(self.transformer.transformer_blocks.attn1.to_q.kernel.value[0].mean()):.6f}")
 
+      # Video-to-Audio Attention weights
+      v2a = self.transformer.transformer_blocks.video_to_audio_attn
+      print(f"WEIGHT DEBUG: block 0 v2a to_q mean: {float(v2a.to_q.kernel.value[0].mean()):.6f}, std: {float(v2a.to_q.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 v2a to_k mean: {float(v2a.to_k.kernel.value[0].mean()):.6f}, std: {float(v2a.to_k.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 v2a to_v mean: {float(v2a.to_v.kernel.value[0].mean()):.6f}, std: {float(v2a.to_v.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 v2a to_out mean: {float(v2a.to_out.kernel.value[0].mean()):.6f}, std: {float(v2a.to_out.kernel.value[0].std()):.6f}")
+
+      # Audio-to-Video Attention weights
+      a2v = self.transformer.transformer_blocks.audio_to_video_attn
+      print(f"WEIGHT DEBUG: block 0 a2v to_q mean: {float(a2v.to_q.kernel.value[0].mean()):.6f}, std: {float(a2v.to_q.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 a2v to_k mean: {float(a2v.to_k.kernel.value[0].mean()):.6f}, std: {float(a2v.to_k.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 a2v to_v mean: {float(a2v.to_v.kernel.value[0].mean()):.6f}, std: {float(a2v.to_v.kernel.value[0].std()):.6f}")
+      print(f"WEIGHT DEBUG: block 0 a2v to_out mean: {float(a2v.to_out.kernel.value[0].mean()):.6f}, std: {float(a2v.to_out.kernel.value[0].std()):.6f}")
+
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
 
       diffusion_loop_start = time.time()