Guard float64 against fp64-less devices in the Helios schedulers

src/diffusers/schedulers/scheduling_helios.py

@@ -22,6 +22,7 @@
 from ..configuration_utils import ConfigMixin, register_to_config
 from ..schedulers.scheduling_utils import SchedulerMixin
 from ..utils import BaseOutput, deprecate
+from ..utils.torch_utils import maybe_adjust_dtype_for_device
 
 
 @dataclass
@@ -235,8 +236,17 @@ def set_timesteps(
             ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
             sigmas = torch.from_numpy(ratios)
 
-        self.timesteps = torch.from_numpy(timesteps).to(device=device)
-        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+        timesteps = torch.from_numpy(timesteps)
+        sigmas = torch.cat([sigmas, torch.zeros(1)])
+        if device is not None:
+            # In the multi-stage branch both arrays come from np.linspace, so they are
+            # float64, which mps (and npu/neuron) cannot hold. Cast before moving.
+            device = torch.device(device)
+            timesteps = timesteps.to(maybe_adjust_dtype_for_device(timesteps.dtype, device))
+            sigmas = sigmas.to(maybe_adjust_dtype_for_device(sigmas.dtype, device))
+
+        self.timesteps = timesteps.to(device=device)
+        self.sigmas = sigmas.to(device=device)
 
         self._step_index = None
         self.reset_scheduler_history()

src/diffusers/schedulers/scheduling_helios_dmd.py

@@ -22,6 +22,7 @@
 from ..configuration_utils import ConfigMixin, register_to_config
 from ..schedulers.scheduling_utils import SchedulerMixin
 from ..utils import BaseOutput
+from ..utils.torch_utils import maybe_adjust_dtype_for_device
 
 
 @dataclass
@@ -213,8 +214,17 @@ def set_timesteps(
             ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
             sigmas = torch.from_numpy(ratios)
 
-        self.timesteps = torch.from_numpy(timesteps).to(device=device)
-        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+        timesteps = torch.from_numpy(timesteps)
+        sigmas = torch.cat([sigmas, torch.zeros(1)])
+        if device is not None:
+            # In the multi-stage branch both arrays come from np.linspace, so they are
+            # float64, which mps (and npu/neuron) cannot hold. Cast before moving.
+            device = torch.device(device)
+            timesteps = timesteps.to(maybe_adjust_dtype_for_device(timesteps.dtype, device))
+            sigmas = sigmas.to(maybe_adjust_dtype_for_device(sigmas.dtype, device))
+
+        self.timesteps = timesteps.to(device=device)
+        self.sigmas = sigmas.to(device=device)
 
         self._step_index = None
         self.reset_scheduler_history()
@@ -275,7 +285,11 @@ def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps):
         # use higher precision for calculations
         original_dtype = flow_pred.dtype
         device = flow_pred.device
-        flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps))
+        # mps (and npu/neuron) cannot hold float64, so fall back to the widest dtype they do
+        dtype = maybe_adjust_dtype_for_device(torch.float64, device)
+        flow_pred, xt, sigmas, timesteps = (
+            x.to(device=device, dtype=dtype) for x in (flow_pred, xt, sigmas, timesteps)
+        )
 
         timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
         sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)

tests/schedulers/test_scheduler_helios.py

@@ -0,0 +1,56 @@
+import unittest
+
+import torch
+
+from diffusers import HeliosDMDScheduler, HeliosScheduler
+from diffusers.utils.torch_utils import maybe_adjust_dtype_for_device
+
+from ..testing_utils import torch_device
+
+
+class HeliosSchedulerDeviceDtypeTest(unittest.TestCase):
+    """Regression tests for #14367.
+
+    Both Helios schedulers built `float64` tensors and moved them straight to the target
+    device. `np.linspace` in the multi-stage branch yields `float64`, so on a device that
+    cannot hold it — mps, and per `_FP64_UNSUPPORTED_DEVICES` also npu and neuron —
+    `set_timesteps` raised `TypeError` before any model ran.
+
+    These assertions are written against `torch_device`, so on a CPU or CUDA runner they
+    check that the dtype is left alone, and on an fp64-less device they check the
+    downcast actually happens.
+    """
+
+    scheduler_classes = (HeliosScheduler, HeliosDMDScheduler)
+    num_inference_steps = 4
+
+    def test_set_timesteps_dtype_is_supported_by_device(self):
+        expected = maybe_adjust_dtype_for_device(torch.float64, torch.device(torch_device))
+        for scheduler_class in self.scheduler_classes:
+            scheduler = scheduler_class()
+            scheduler.set_timesteps(self.num_inference_steps, device=torch_device, stage_index=0)
+
+            for name, tensor in (("timesteps", scheduler.timesteps), ("sigmas", scheduler.sigmas)):
+                assert tensor.device.type == torch.device(torch_device).type, (
+                    f"{scheduler_class.__name__}.{name} was not moved to {torch_device}"
+                )
+                assert tensor.dtype == expected, (
+                    f"{scheduler_class.__name__}.{name} has dtype {tensor.dtype} on {torch_device}, "
+                    f"expected {expected}"
+                )
+
+    def test_convert_flow_pred_to_x0_runs_on_device(self):
+        scheduler = HeliosDMDScheduler()
+        scheduler.set_timesteps(self.num_inference_steps, device=torch_device, stage_index=0)
+
+        shape = (1, 4, 2, 8, 8)
+        flow_pred = torch.randn(shape, generator=torch.Generator().manual_seed(0)).to(torch_device)
+        xt = torch.randn(shape, generator=torch.Generator().manual_seed(1)).to(torch_device)
+
+        x0_pred = scheduler.convert_flow_pred_to_x0(
+            flow_pred, xt, scheduler.timesteps[:1], scheduler.sigmas, scheduler.timesteps
+        )
+
+        assert x0_pred.shape == flow_pred.shape
+        assert x0_pred.dtype == flow_pred.dtype
+        assert torch.isfinite(x0_pred).all()