diff --git a/src/diffusers/schedulers/scheduling_helios.py b/src/diffusers/schedulers/scheduling_helios.py index 6d24e54627aa..7bd18b5ad38d 100644 --- a/src/diffusers/schedulers/scheduling_helios.py +++ b/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() diff --git a/src/diffusers/schedulers/scheduling_helios_dmd.py b/src/diffusers/schedulers/scheduling_helios_dmd.py index 1f4afa0e3128..661016a2ed51 100644 --- a/src/diffusers/schedulers/scheduling_helios_dmd.py +++ b/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) diff --git a/tests/schedulers/test_scheduler_helios.py b/tests/schedulers/test_scheduler_helios.py new file mode 100644 index 000000000000..48a84e93217a --- /dev/null +++ b/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()