Guard float64 against fp64-less devices in the Helios schedulers

[4ktLuffy]

What this fixes

Fixes #14367.

Both Helios schedulers build float64 tensors and move them straight to the target device. In the multi-stage branch, timesteps and sigmas both come from np.linspace, which is float64, so on a device that cannot hold it the move raises:

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64.

HeliosPyramidPipeline calls set_timesteps(..., device=device), so on Apple Silicon this fails at the first scheduler call, before any model runs.

Three sites, all reproducible on main (0.40.0.dev0):

Site	Scheduler
scheduling_helios.py set_timesteps	HeliosScheduler
scheduling_helios_dmd.py set_timesteps	HeliosDMDScheduler
scheduling_helios_dmd.py convert_flow_pred_to_x0 — .double() on four on-device tensors	HeliosDMDScheduler

The change

All three route the dtype through maybe_adjust_dtype_for_device, which this codebase already uses for exactly this (53 call sites). It downcasts float64 only on devices in _FP64_UNSUPPORTED_DEVICES — mps, npu, neuron — and is a no-op elsewhere, so CPU and CUDA keep full float64 precision.

In set_timesteps the cast happens before the move rather than after, since the failing operation is the move.

Verification

	before	after
HeliosScheduler.set_timesteps(device="mps")	TypeError	ok, float32
HeliosDMDScheduler.set_timesteps(device="mps")	TypeError	ok, float32
HeliosDMDScheduler.convert_flow_pred_to_x0 on mps	TypeError	ok, finite
the same three on CPU	ok, float64	bit-identical, float64

Full tests/schedulers/: 6 failed / 979 passed with this change, versus 8 failed / 977 passed without it — the same six pre-existing local MPS failures either way, plus the two new tests.

Tests

tests/schedulers/test_scheduler_helios.py is new; these schedulers had no test coverage. Both tests fail on main and pass with this change.

They assert against torch_device, comparing the resulting dtype to maybe_adjust_dtype_for_device(torch.float64, torch_device). On a CPU or CUDA runner that asserts float64 is left alone; on an fp64-less device it asserts the downcast happened. So the tests are meaningful on CI without requiring Apple hardware, though only an mps/npu/neuron runner exercises the actual regression.

Overlap with #14071

@Shreyas-jk opened #14071 in June, which fixes the third site (convert_flow_pred_to_x0). I did not spot it when I filed #14367 — apologies for the duplication, that was my miss.

The two differences, so you can pick rather than guess:

• Coverage. [mps] Fix float64 upcast in HeliosDMDScheduler.convert_flow_pred_to_x0 #14071 fixes the convert_flow_pred_to_x0 upcast only. The two set_timesteps sites are untouched by it, and those are the ones that break HeliosPyramidPipeline first — it calls set_timesteps(..., device=device) before any model runs, so on Apple Silicon the failure happens there, not in convert_flow_pred_to_x0.
• Mechanism. [mps] Fix float64 upcast in HeliosDMDScheduler.convert_flow_pred_to_x0 #14071 uses torch.float32 if device.type == "mps" else torch.float64. Routing through maybe_adjust_dtype_for_device instead picks up npu and neuron from _FP64_UNSUPPORTED_DEVICES as well, and keeps this file consistent with the other 53 call sites.

I am happy to drop my third hunk and rebase on top of #14071 if you would rather merge that one first — its author got there before me and the credit for that site is theirs either way. Just say which you prefer.

[Shreyas-jk]

@4ktLuffy — sorry for the slow reply, and thanks for flagging the overlap rather than quietly shipping past it.

I checked both of your points against the tree and agree yours should land. Closing #14071 in favour of this.

On coverage: you're right that set_timesteps is the site that matters and mine doesn't touch it — both schedulers still have the bare torch.from_numpy(timesteps).to(device=device) at scheduling_helios.py:238 and scheduling_helios_dmd.py:216. One small correction: in HeliosPyramidPipeline the failure isn't before any model runs — encode_prompt (:710) and prepare_latents (:887) both execute before set_timesteps at :941. It fails before the transformer, not before the text encoder or VAE. Doesn't change the conclusion.

On mechanism: agreed, and it's the stronger reason. maybe_adjust_dtype_for_device picks up npu and neuron via _FP64_UNSUPPORTED_DEVICES, where my device.type == "mps" check only ever covered one. Routing through the existing helper is the right call.

Your tests are also better positioned than mine — asserting against torch_device means they're meaningful on the CPU/CUDA runners, where mine was skipUnless(mps) and never ran in CI. No need to preserve anything from my PR.

No need to drop your third hunk or rebase on me. Ship it as-is.