From ad07b06de59d6290c20c4d4b26f68c237c9df3e1 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 10 Jul 2026 10:21:34 -0600 Subject: [PATCH] Use cached first frame for wan22_5 model --- .../diffusion_models/wan22/wan22_5b_model.py | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py index 5d8e48f..b3aed16 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py @@ -265,23 +265,37 @@ class Wan225bModel(Wan21): if batch.dataset_config.do_i2v: with torch.no_grad(): - frames = batch.tensor - if len(frames.shape) == 4: - first_frames = frames - elif len(frames.shape) == 5: - first_frames = frames[:, 0] - # Add conditioning using the standalone function + # check to see if we had the first frame latent cached + if batch.first_frame_latents is not None: conditioned_latent, noise_mask = add_first_frame_conditioning_v22( latent_model_input=latent_model_input.to( self.device_torch, self.torch_dtype ), - first_frame=first_frames.to(self.device_torch, self.torch_dtype), + first_frame_latents=batch.first_frame_latents.to( + self.device_torch, self.torch_dtype + ), vae=self.vae, ) # conditioned tokens are clean with timestep 0 and must not contribute to the loss self._i2v_loss_mask = noise_mask else: - raise ValueError(f"Unknown frame shape {frames.shape}") + frames = batch.tensor + if len(frames.shape) == 4: + first_frames = frames + elif len(frames.shape) == 5: + first_frames = frames[:, 0] + # Add conditioning using the standalone function + conditioned_latent, noise_mask = add_first_frame_conditioning_v22( + latent_model_input=latent_model_input.to( + self.device_torch, self.torch_dtype + ), + first_frame=first_frames.to(self.device_torch, self.torch_dtype), + vae=self.vae, + ) + # conditioned tokens are clean with timestep 0 and must not contribute to the loss + self._i2v_loss_mask = noise_mask + else: + raise ValueError(f"Unknown frame shape {frames.shape}") # make the noise mask if noise_mask is None: