Added more debug for NaNs

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -373,7 +373,6 @@ def print_step_zero_stats(operands):
                     std=jnp.std(latents),
                     mean=jnp.mean(latents))
     jax.debug.print("first_frame_mask shape: {}", first_frame_mask.shape if first_frame_mask is not None else (-1,))
-    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()