changed expand_timesteps logic

prishajain1 · prishajain1 · commit 390e11321f8b · 2025-12-30T14:49: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
@@ -255,7 +255,8 @@ def __call__(
       )
       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
+         clean_latents = condition[..., 4:]
+         latents = first_frame_mask * clean_latents + (1 - 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
`@@ -255,7 +255,8 @@ def __call__(`
`255`	`255`	`)`
`256`	`256`	`if self.config.expand_timesteps:`
`257`	`257`	`jax.debug.print("Applying first frame preservation with expand_timesteps.")`
`258`		`- latents = (1 - first_frame_mask) * condition + first_frame_mask * latents`
	`258`	`+ clean_latents = condition[..., 4:]`
	`259`	`+ latents = first_frame_mask * clean_latents + (1 - first_frame_mask) * latents`
`259`	`260`	`latents_bcthw = jnp.transpose(latents, (0, 4, 1, 2, 3))`
`260`	`261`	`latents_denorm_bcthw = self._denormalize_latents(latents_bcthw)`
`261`	`262`