huggingface / huggingface/diffusers
Helios schedulers raise `TypeError` on MPS: `float64` without the device guard
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
### 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:
```python
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:
```python
# 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
Beitragsleitfaden
Rechercherichtung
Start with scheduling_helios.py and scheduling_helios_dmd.py at the reported set_timesteps and convert_flow_pred_to_x0 sites, then read maybe_adjust_dtype_for_device in utils/torch_utils.py and the sibling handling in scheduling_flow_map_euler_discrete.py. Run the issue's CPU and MPS scheduler reproductions, including HeliosDMDScheduler.step(). Done means both schedulers avoid unsupported float64 operations on MPS while preserving higher-precision calculations where the device supports them.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- machine-learning
- Issue-Typ
- Bug
- Schwierigkeit
- 3/5
- Geschätzter Aufwand
- 1-2 Tage
- Aktivitätsstatus
- Aktiv
- Klarheit
- Klar beschrieben
- Anfängerfreundlichkeit
- 74/100