more debug added

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -152,6 +152,7 @@ def prepare_latents(
         jax.debug.print("condition stats: mask_mean={mm}, latent_mean={lm}",
                         mm=jnp.mean(condition[..., 0]),
                         lm=jnp.mean(condition[..., 1:]))
+        jax.debug.print("condition latent std={std}", std=jnp.std(condition[..., 1:]))
 
         return latents, condition, None
 
@@ -303,15 +304,31 @@ def loop_body(step, vals):
     original_dtype = latents.dtype
     rng, timestep_rng = jax.random.split(rng)
     t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
+    jax.debug.print("Step {s}: timestep={t}", s=step, t=t)
 
     latents_input = latents
     if do_classifier_free_guidance:
         latents_input = jnp.concatenate([latents, latents], axis=0)
+    jax.debug.print("Step{s}: latents_input stats min={mn}, max={mx}, mean={mean}, std={std}",
+                    s=step,
+                    mn=jnp.min(latents_input),
+                    mx=jnp.max(latents_input),
+                    mean=jnp.mean(latents_input),
+                    std=jnp.std(latents_input))
 
     latent_model_input = jnp.concatenate([latents_input, condition], axis=-1)
     timestep = jnp.broadcast_to(t, latents_input.shape[0])
     latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
 
+    jax.debug.print("Step {s}: latent_model_input shape: {shape}",
+                    s=step,
+                    shape=latent_model_input.shape)
+
+    channel_energy = jnp.sum(latent_model_input*latent_model_input,axis=(0,2,3,4))
+    jax.debug.print("Step {s}: channel energy first 10={ce}",
+                    s=step,
+                    ce=channel_energy[:10])
+
     prompt_embeds_input = prompt_embeds
     image_embeds_input = image_embeds
 
@@ -324,10 +341,20 @@ def loop_body(step, vals):
         encoder_hidden_states_image=image_embeds_input,
     )
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
+    jax.debug.print("Step {s}: noise_pred stats min={mn}, max={mx}, mean={mean}, std={std}",
+                    s=step,
+                    mn=jnp.min(noise_pred),
+                    mx=jnp.max(noise_pred),
+                    mean=jnp.mean(noise_pred),
+                    std=jnp.std(noise_pred))
     jax.debug.print("Step {s}: latents_prev std={std}, mean={mean}",
                     s=step,
                     std=jnp.std(latents),
                     mean=jnp.mean(latents))
+    jax.debug.print("first_frame_mask shape:", first_frame_mask.shape)
+    jax.debug.print("first_frame_mask unique values:", jnp.unique(first_frame_mask))
+    jax.debug.print("condition shape:", condition.shape)
+    jax.debug.print("condition stats:", jnp.min(condition), jnp.max(condition), jnp.mean(condition))
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
     jax.debug.print("Step {s}: latents_next std={std}, mean={mean}",
                     s=step,