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.