checking expand timesteps

prishajain1 · prishajain1 · commit bc6955d39eb3 · 2025-12-30T14:47:12.000+05:30
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py
@@ -254,6 +254,7 @@ def __call__(
           rng=inference_rng,
       )
       if self.config.expand_timesteps:
+         jax.debug.print("Applying first frame preservation with expand_timesteps.")
          latents = (1 - first_frame_mask) * condition + first_frame_mask * latents
       latents_bcthw = jnp.transpose(latents, (0, 4, 1, 2, 3))
       latents_denorm_bcthw = self._denormalize_latents(latents_bcthw)

Original file line number	Diff line number	Diff line change
`@@ -254,6 +254,7 @@ def __call__(`
`254`	`254`	`rng=inference_rng,`
`255`	`255`	`)`
`256`	`256`	`if self.config.expand_timesteps:`
	`257`	`+ jax.debug.print("Applying first frame preservation with expand_timesteps.")`
`257`	`258`	`latents = (1 - first_frame_mask) * condition + first_frame_mask * latents`
`258`	`259`	`latents_bcthw = jnp.transpose(latents, (0, 4, 1, 2, 3))`
`259`	`260`	`latents_denorm_bcthw = self._denormalize_latents(latents_bcthw)`