Fix Helios auto offload decode

[pzarzycki]

Summary

Fixes HeliosDecodeStep when modular Helios pipelines run with component auto CPU offload enabled.

The decode step was normalizing chunk latents using latents_mean / latents_std tensors created on the VAE device, while the chunk latents themselves could remain on a different device under offload. That caused device mismatches during decode.

This change aligns Helios with the working Wan decoder pattern by normalizing each chunk on the chunk latent tensor's own device and dtype before converting to the VAE decode dtype.

Repro

python -m pytest tests/modular_pipelines/helios/test_modular_pipeline_helios.py::TestHeliosModularPipelineFast::test_components_auto_cpu_offload_inference_consistent -q

Before this change, the test failed in src/diffusers/modular_pipelines/helios/decoders.py with a device mismatch during decode under auto CPU offload.

Implementation

• updated HeliosDecodeStep in src/diffusers/modular_pipelines/helios/decoders.py
• compute latents_mean / latents_std on each chunk latent tensor's device and dtype
• normalize chunk latents there first, then cast to the VAE decode dtype
• keep the rest of the decode flow unchanged

Tests

Ran on accelerator-backed setup:

python -m pytest tests/modular_pipelines/helios/test_modular_pipeline_helios.py::TestHeliosModularPipelineFast::test_components_auto_cpu_offload_inference_consistent -q
python -m pytest tests/modular_pipelines/helios/test_modular_pipeline_helios.py -q

Results:

1 passed
24 passed, 4 skipped

Also ran CPU-mode validation:

DIFFUSERS_TEST_DEVICE=cpu python -m pytest tests/modular_pipelines/helios/test_modular_pipeline_helios.py -q

Result:

17 passed, 11 skipped

[Copilot]

Pull request overview

Fixes a device/dtype mismatch in HeliosDecodeStep that occurs when modular Helios pipelines run with component auto CPU offload enabled, by normalizing latents on the per-chunk latent tensor’s device/dtype before casting to the VAE decode dtype.

Changes:

• Move latents_mean / latents_std construction to use each chunk_latents tensor’s device and dtype.
• Normalize each chunk in-place on the chunk’s device/dtype, then cast to vae.dtype prior to vae.decode.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yiyixuxu]

Can we just use this?

device = components._exeution_device

[pzarzycki]

I updated this to use components._execution_device in HeliosDecodeStep and moved the latents_mean / latents_std setup out of the chunk loop.

Using chunk_latents.device fixes the mismatch, but it also rebuilds those tensors for every chunk. Since the modular pipeline already uses components._execution_device as the execution source of truth in nearby Helios steps, that seemed like the cleaner fix here.

I also added a regression test for the offload case where vae.device is CPU but decode runs on the execution device.

I reran the full Helios modular test file after that:

• DIFFUSERS_TEST_DEVICE=cpu .../test_modular_pipeline_helios.py -> 17 passed, 12 skipped
• DIFFUSERS_TEST_DEVICE=cuda .../test_modular_pipeline_helios.py -> 25 passed, 4 skipped

[github-actions]

Hi @pzarzycki, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

can you remove the test? we already have offloading test https://github.com/huggingface/diffusers/blob/main/tests/modular_pipelines/test_modular_pipelines_common.py#L350

[pzarzycki]

Removed the Helios-specific offload regression test as suggested.

I kept the decoder fix and dropped the redundant test/helper block in tests/modular_pipelines/helios/test_modular_pipeline_helios.py, since the shared modular offloading coverage in tests/modular_pipelines/test_modular_pipelines_common.py already covers that behavior.

I also fixed the pending formatter issue on src/diffusers/modular_pipelines/helios/decoders.py so the branch passes the quality check.

Local validation:

• pytest tests/modular_pipelines/helios/test_modular_pipeline_helios.py -q -> 24 passed, 4 skipped
• ruff check examples scripts src tests utils benchmarks setup.py
• ruff format --check examples scripts src tests utils benchmarks setup.py

[yiyixuxu]

thanks