From 357372a73318b5bfd9ee7ea92bcd36fa95753570 Mon Sep 17 00:00:00 2001 From: Pawel Date: Tue, 7 Jul 2026 20:40:18 -0700 Subject: [PATCH 1/3] Fix Helios auto offload decode --- .../modular_pipelines/helios/decoders.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/diffusers/modular_pipelines/helios/decoders.py b/src/diffusers/modular_pipelines/helios/decoders.py index 4f7328c32c7d..2e488f170f4c 100644 --- a/src/diffusers/modular_pipelines/helios/decoders.py +++ b/src/diffusers/modular_pipelines/helios/decoders.py @@ -77,16 +77,18 @@ def __call__(self, components, state: PipelineState) -> PipelineState: vae = components.vae - latents_mean = ( - torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.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 - ) - history_video = None for chunk_latents in block_state.latent_chunks: - current_latents = chunk_latents.to(vae.dtype) / latents_std + latents_mean + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to( + chunk_latents.device, chunk_latents.dtype + ) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + chunk_latents.device, chunk_latents.dtype + ) + current_latents = chunk_latents / latents_std + latents_mean + current_latents = current_latents.to(vae.dtype) current_video = vae.decode(current_latents, return_dict=False)[0] if history_video is None: From 6323bdd4492a38c28920e32b6fedb32e4fbadf38 Mon Sep 17 00:00:00 2001 From: Pawel Date: Wed, 8 Jul 2026 19:31:17 -0700 Subject: [PATCH 2/3] Fix Helios decode offload device handling and add regression test --- .../modular_pipelines/helios/decoders.py | 18 +++---- .../helios/test_modular_pipeline_helios.py | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/src/diffusers/modular_pipelines/helios/decoders.py b/src/diffusers/modular_pipelines/helios/decoders.py index 2e488f170f4c..e83f965571d8 100644 --- a/src/diffusers/modular_pipelines/helios/decoders.py +++ b/src/diffusers/modular_pipelines/helios/decoders.py @@ -76,19 +76,17 @@ 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(device, decode_dtype) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + device, decode_dtype + ) history_video = None for chunk_latents in block_state.latent_chunks: - latents_mean = ( - torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to( - chunk_latents.device, chunk_latents.dtype - ) - ) - latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( - chunk_latents.device, chunk_latents.dtype - ) - current_latents = chunk_latents / latents_std + latents_mean - current_latents = current_latents.to(vae.dtype) + 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: diff --git a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py index 44a01dad6525..a42df56e9aa1 100644 --- a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py +++ b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py @@ -13,7 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace + import pytest +import torch from diffusers.modular_pipelines import ( HeliosAutoBlocks, @@ -21,7 +24,10 @@ HeliosPyramidAutoBlocks, HeliosPyramidModularPipeline, ) +from diffusers.modular_pipelines.helios.decoders import HeliosDecodeStep +from diffusers.modular_pipelines.modular_pipeline import PipelineState +from ...testing_utils import torch_device from ..test_modular_pipelines_common import ModularPipelineTesterMixin @@ -91,6 +97,51 @@ def test_num_images_per_prompt(self): pass +class DummyHeliosVAE: + def __init__(self): + self.config = SimpleNamespace(latents_mean=[0.5, -0.5], latents_std=[2.0, 4.0], z_dim=2) + self.dtype = torch.float32 + self.device = torch.device("cpu") + self.last_decode_latents = None + + def decode(self, latents, return_dict=False): + self.last_decode_latents = latents + return (latents,) + + +class DummyVideoProcessor: + def postprocess_video(self, history_video, output_type="pt"): + return history_video + + +@pytest.mark.skipif(torch_device == "cpu", reason="Helios offload regression requires an accelerator device") +def test_helios_decode_uses_execution_device_for_chunk_latents(): + decode_step = HeliosDecodeStep() + components = SimpleNamespace( + vae=DummyHeliosVAE(), + video_processor=DummyVideoProcessor(), + vae_scale_factor_temporal=4, + _execution_device=torch.device(torch_device), + ) + latent_chunks = [ + torch.arange(18, device=torch_device, dtype=torch.float32).reshape(1, 2, 9, 1, 1), + ] + state = PipelineState() + state.set("latent_chunks", latent_chunks) + state.set("num_frames", 33) + state.set("output_type", "pt") + + _, state = decode_step(components, state) + + expected = latent_chunks[0].to(device=torch_device, dtype=torch.float32) + expected = expected / torch.tensor([0.5, 0.25], device=torch_device, dtype=torch.float32).view(1, 2, 1, 1, 1) + expected = expected + torch.tensor([0.5, -0.5], device=torch_device, dtype=torch.float32).view(1, 2, 1, 1, 1) + + torch.testing.assert_close(components.vae.last_decode_latents, expected) + assert components.vae.last_decode_latents.device.type == torch_device + torch.testing.assert_close(state.videos, expected) + + HELIOS_PYRAMID_WORKFLOWS = { "text2video": [ ("text_encoder", "HeliosTextEncoderStep"), From 66faa96d73f9e6298f0ac306888e3256dec63315 Mon Sep 17 00:00:00 2001 From: Pawel Date: Fri, 10 Jul 2026 08:37:31 -0700 Subject: [PATCH 3/3] Remove redundant Helios offload test --- .../modular_pipelines/helios/decoders.py | 4 +- .../helios/test_modular_pipeline_helios.py | 51 ------------------- 2 files changed, 3 insertions(+), 52 deletions(-) diff --git a/src/diffusers/modular_pipelines/helios/decoders.py b/src/diffusers/modular_pipelines/helios/decoders.py index e83f965571d8..c448d36136e6 100644 --- a/src/diffusers/modular_pipelines/helios/decoders.py +++ b/src/diffusers/modular_pipelines/helios/decoders.py @@ -79,7 +79,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: 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(device, decode_dtype) + latents_mean = ( + 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( device, decode_dtype ) diff --git a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py index a42df56e9aa1..44a01dad6525 100644 --- a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py +++ b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py @@ -13,10 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from types import SimpleNamespace - import pytest -import torch from diffusers.modular_pipelines import ( HeliosAutoBlocks, @@ -24,10 +21,7 @@ HeliosPyramidAutoBlocks, HeliosPyramidModularPipeline, ) -from diffusers.modular_pipelines.helios.decoders import HeliosDecodeStep -from diffusers.modular_pipelines.modular_pipeline import PipelineState -from ...testing_utils import torch_device from ..test_modular_pipelines_common import ModularPipelineTesterMixin @@ -97,51 +91,6 @@ def test_num_images_per_prompt(self): pass -class DummyHeliosVAE: - def __init__(self): - self.config = SimpleNamespace(latents_mean=[0.5, -0.5], latents_std=[2.0, 4.0], z_dim=2) - self.dtype = torch.float32 - self.device = torch.device("cpu") - self.last_decode_latents = None - - def decode(self, latents, return_dict=False): - self.last_decode_latents = latents - return (latents,) - - -class DummyVideoProcessor: - def postprocess_video(self, history_video, output_type="pt"): - return history_video - - -@pytest.mark.skipif(torch_device == "cpu", reason="Helios offload regression requires an accelerator device") -def test_helios_decode_uses_execution_device_for_chunk_latents(): - decode_step = HeliosDecodeStep() - components = SimpleNamespace( - vae=DummyHeliosVAE(), - video_processor=DummyVideoProcessor(), - vae_scale_factor_temporal=4, - _execution_device=torch.device(torch_device), - ) - latent_chunks = [ - torch.arange(18, device=torch_device, dtype=torch.float32).reshape(1, 2, 9, 1, 1), - ] - state = PipelineState() - state.set("latent_chunks", latent_chunks) - state.set("num_frames", 33) - state.set("output_type", "pt") - - _, state = decode_step(components, state) - - expected = latent_chunks[0].to(device=torch_device, dtype=torch.float32) - expected = expected / torch.tensor([0.5, 0.25], device=torch_device, dtype=torch.float32).view(1, 2, 1, 1, 1) - expected = expected + torch.tensor([0.5, -0.5], device=torch_device, dtype=torch.float32).view(1, 2, 1, 1, 1) - - torch.testing.assert_close(components.vae.last_decode_latents, expected) - assert components.vae.last_decode_latents.device.type == torch_device - torch.testing.assert_close(state.videos, expected) - - HELIOS_PYRAMID_WORKFLOWS = { "text2video": [ ("text_encoder", "HeliosTextEncoderStep"),