diff --git a/src/diffusers/modular_pipelines/helios/decoders.py b/src/diffusers/modular_pipelines/helios/decoders.py index 4f7328c32c7d..c448d36136e6 100644 --- a/src/diffusers/modular_pipelines/helios/decoders.py +++ b/src/diffusers/modular_pipelines/helios/decoders.py @@ -76,17 +76,19 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) vae = components.vae + device = components._execution_device + decode_dtype = vae.dtype latents_mean = ( - torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.dtype) + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(device, decode_dtype) ) latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( - vae.device, vae.dtype + device, decode_dtype ) history_video = None for chunk_latents in block_state.latent_chunks: - current_latents = chunk_latents.to(vae.dtype) / latents_std + latents_mean + current_latents = chunk_latents.to(device=device, dtype=decode_dtype) / latents_std + latents_mean current_video = vae.decode(current_latents, return_dict=False)[0] if history_video is None: