Helios schedulers raise `TypeError` on MPS: `float64` without the device guard

[4ktLuffy]

Describe the bug

Both Helios schedulers raise TypeError on MPS (Apple Silicon) because they build float64 tensors without the device guard the rest of the codebase uses. HeliosPyramidPipeline calls set_timesteps(..., device=device), so on an MPS device the pipeline fails at the first scheduler call, before any model runs.

Two sites, both verified by traceback on main (0.40.0.dev0):

1. set_timesteps — affects both schedulers

scheduling_helios.py:238      self.timesteps = torch.from_numpy(timesteps).to(device=device)
scheduling_helios_dmd.py:216  self.timesteps = torch.from_numpy(timesteps).to(device=device)

timesteps is a NumPy float64 array, so torch.from_numpy yields float64 and .to(device="mps") raises.

2. convert_flow_pred_to_x0 — HeliosDMDScheduler only, reached from step()

scheduling_helios_dmd.py:278  flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (...))

Reproduction

CPU only for the passing case, MPS for the failing one — no model weights needed:

import torch
from diffusers import HeliosScheduler, HeliosDMDScheduler

# 1. set_timesteps on an MPS device — both schedulers
for name, cls in (("HeliosScheduler", HeliosScheduler), ("HeliosDMDScheduler", HeliosDMDScheduler)):
    try:
        cls().set_timesteps(4, device="mps", stage_index=0)
        print(name, "ok")
    except TypeError as e:
        print(name, "TypeError:", str(e)[:60])

# 2. HeliosDMDScheduler.step() with tensors already on MPS
d = HeliosDMDScheduler()
a = torch.randn(1, 4, 2, 8, 8, device="mps")
b = torch.randn(1, 4, 2, 8, 8, device="mps")
ts = torch.tensor([1000., 750., 500., 250.], device="mps")
sg = torch.tensor([1., .75, .5, .25], device="mps")
d.step(model_output=a, timestep=1000.0, sample=b, cur_sampling_step=3,
       dmd_noisy_tensor=torch.zeros_like(b), dmd_sigmas=sg,
       dmd_timesteps=ts, all_timesteps=ts)

HeliosScheduler TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS
HeliosDMDScheduler TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS

Traceback (most recent call last):
  scheduling_helios_dmd.py:298 in step
  scheduling_helios_dmd.py:278 in convert_flow_pred_to_x0
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64.

Both schedulers work correctly on CPU with the same inputs.

Expected behavior

The intent at scheduling_helios_dmd.py:277 is stated as "use higher precision for calculations" — the goal is precision, not float64 specifically, so a downcast on backends that lack it preserves the intent.

The codebase already has the mechanism. maybe_adjust_dtype_for_device in utils/torch_utils.py maps float64 -> float32 for mps/npu/neuron and is used at ~53 call sites. Neither Helios scheduler imports it. A sibling scheduler handles the same problem explicitly:

# scheduling_flow_map_euler_discrete.py:169-173
# ...the final tensors to the requested device (with a float32 downcast for MPS / NPU).
is_mps = device_obj is not None and device_obj.type == "mps"
is_npu = device_obj is not None and device_obj.type == "npu"
out_dtype = torch.float32 if (is_mps or is_npu) else torch.float64

Scope

Verified at the scheduler level only. I have not run HeliosPyramidPipeline end to end on MPS — that needs the Helios-14B weights, which I cannot download here — so I have not established whether the pipeline would otherwise work on MPS once the schedulers are fixed. What is established is that pipeline_helios_pyramid.py:941 passes device=device into set_timesteps, so this is on the pipeline's path rather than an artificial call.

Also not tested: npu and neuron, which _DTYPE_UNSUPPORTED_DEVICES lists alongside mps for float64. They may be affected identically.

Unrelated to #14353, which covers HeliosDMDScheduler's ignored config options in the same file — different defect, no overlap in the fix.

System Info

• diffusers version: 0.40.0.dev0 (main, commit a8345366e)
• Platform: macOS 26.0.1, arm64 (Apple Silicon)
• Python version: 3.11
• PyTorch version: 2.13.0, torch.backends.mps.is_available() == True

Who can help?

@dg845 (git blame points to #13208 for both files) and @yiyixuxu