debug added in wan 2.1 i2v pipeline

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -256,6 +256,16 @@ def __call__(
       max_logging.log(f"[DEBUG CALL] Decoded video type: {type(decoded_video)}")
     return decoded_video
 
+def check_nan_jit(tensor: jax.Array, name: str, step: jax.Array):
+    if tensor is None:
+      return
+
+    has_nans = jnp.isnan(tensor).any()
+    has_infs = jnp.isinf(tensor).any()
+    jax.debug.print(f"[DEBUG JIT {jax.process_index()}] Step: {{step}} - {name}: "
+                    "Shape: {shape}, Has NaNs: {has_nans_val}, Has Infs: {has_infs_val}",
+                    step=step, shape=tensor.shape, has_nans_val=has_nans, has_infs_val=has_infs)
+
 def run_inference_2_1_i2v(
     graphdef, sharded_state, rest_of_state,
     latents: jnp.array,
@@ -289,16 +299,24 @@ def loop_body(step, vals):
     rng, timestep_rng = jax.random.split(rng)
     t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
 
+    check_nan_jit(latents, "latents_prev at loop start", step)
+
     latents_input = latents
     if do_classifier_free_guidance:
         latents_input = jnp.concatenate([latents, latents], axis=0)
+    check_nan_jit(latents_input, "latents_input after CFG concat", step)
 
     latent_model_input = jnp.concatenate([latents_input, condition], axis=-1)
+    check_nan_jit(latent_model_input, "latent_model_input after cond concat", step)
     timestep = jnp.broadcast_to(t, latents_input.shape[0])
     latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
+    check_nan_jit(latent_model_input, "latent_model_input for transformer", step)
 
     prompt_embeds_input = prompt_embeds
     image_embeds_input = image_embeds
+    check_nan_jit(prompt_embeds_input, "prompt_embeds_input for transformer", step)
+    check_nan_jit(image_embeds_input, "image_embeds_input for transformer", step)
+
 
     noise_pred, _ = transformer_forward_pass(
         graphdef, sharded_state, rest_of_state,
@@ -307,11 +325,17 @@ def loop_body(step, vals):
         guidance_scale=guidance_scale,
         encoder_hidden_states_image=image_embeds_input,
     )
+    check_nan_jit(noise_pred, "noise_pred_bcthw from transformer", step)
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
+    check_nan_jit(noise_pred, "noise_pred after transpose", step)
 
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
+    check_nan_jit(latents, "latents_next after scheduler", step)
     latents = latents.astype(original_dtype)
+    check_nan_jit(latents, "latents_next after dtype cast", step)
     return latents, scheduler_state, rng
 
+  max_logging.log(f"Running fori_loop for {num_inference_steps} steps.")
   latents, _, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state, rng))
+  max_logging.log("Finished fori_loop.")
   return latents