diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py index f836df424ad0..0f2291d58fd7 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py @@ -323,8 +323,12 @@ def set_timesteps( sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) sigmas = torch.from_numpy(sigmas).to(device=device) - # interpolate sigmas - sigmas_interpol = sigmas.log().lerp(sigmas.roll(1).log(), 0.5).exp() + # interpolate sigmas: the geometric mean of each sigma and its predecessor. + # Computed directly rather than as exp(lerp(log, log)), because the schedule always + # ends in a zero sigma and log(0) = -inf makes that lerp return NaN. torch resolves + # the -inf differently on CPU and MPS (pytorch#111374), so the NaN landed on a dead + # entry on CPU but on live entries on MPS, where it propagated into the sample. + sigmas_interpol = (sigmas * sigmas.roll(1)).sqrt() self.sigmas = torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2), sigmas[-1:]]) self.sigmas_interpol = torch.cat( diff --git a/tests/schedulers/test_scheduler_kdpm2_discrete.py b/tests/schedulers/test_scheduler_kdpm2_discrete.py index 370ba2253ee2..ae23d95f0149 100644 --- a/tests/schedulers/test_scheduler_kdpm2_discrete.py +++ b/tests/schedulers/test_scheduler_kdpm2_discrete.py @@ -170,3 +170,22 @@ def test_beta_sigmas(self): def test_exponential_sigmas(self): self.check_over_configs(use_exponential_sigmas=True) + + def test_set_timesteps_sigmas_interpol_no_nan(self): + # Regression test for #14368: sigmas_interpol was computed as + # exp(lerp(log(sigmas), log(sigmas.roll(1)), 0.5)). The schedule always ends in a + # zero sigma, so log(0) = -inf entered the lerp, and torch resolves that -inf + # differently on CPU and MPS (pytorch#111374) — leaving a NaN on a dead entry on + # CPU but on live entries on MPS, where it propagated into every sample. + scheduler_class = self.scheduler_classes[0] + for num_inference_steps in (4, 10, 25): + scheduler = scheduler_class(**self.get_scheduler_config()) + scheduler.set_timesteps(num_inference_steps, device=torch_device) + + assert torch.isfinite(scheduler.sigmas_interpol).all(), ( + f"set_timesteps({num_inference_steps}) produced non-finite sigmas_interpol " + f"on {torch_device}: {scheduler.sigmas_interpol}" + ) + assert torch.isfinite(scheduler.sigmas).all(), ( + f"set_timesteps({num_inference_steps}) produced non-finite sigmas on {torch_device}" + )