From 57825236890ca96019a0d0e0c6b7d185cab0c8be Mon Sep 17 00:00:00 2001 From: 4ktluffy Date: Mon, 3 Aug 2026 20:42:36 +0300 Subject: [PATCH] Fix NaN in KDPM2DiscreteScheduler.sigmas_interpol MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sigmas_interpol was computed as exp(lerp(log(sigmas), log(sigmas.roll(1)), 0.5)), the geometric mean of each sigma and its predecessor taken through log space. The sigma schedule always ends in a zero sigma, so log(0) = -inf entered the lerp. torch resolves that -inf differently on CPU and MPS (pytorch#111374), so: - on CPU the NaN landed on sigmas_interpol[0], a wrap-around entry that step() never reads — the value was wrong but harmless, and present in every config tested (24/24); - on MPS it landed on the trailing live entries instead, propagating through dt = sigma_interpol - sigma_hat into prev_sample and turning the whole latent NaN. A tiny StableDiffusionPipeline run with this scheduler returns 49152/49152 NaN pixels — a fully black image. Computing the geometric mean directly as sqrt(a * b) avoids log(0) entirely and is exactly equivalent for positive sigmas. Measured against a float64 ground truth it is also slightly more accurate than the log/exp round trip: max relative error 5.0e-08 vs 1.4e-07, better on 8 of 25 entries and worse on none. The largest sigma across configs is 2.0e+04, so the squared intermediate has ample float32 headroom. After the change MPS output matches CPU exactly (max delta 0.0 across 15 configurations) and sigmas_interpol is finite on both devices in all 48 configurations tested. The ancestral variant has the same log/lerp pattern but is left alone here: #14221 is already rewriting that line as part of the fix for #14213. --- .../schedulers/scheduling_k_dpm_2_discrete.py | 8 ++++++-- .../test_scheduler_kdpm2_discrete.py | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) 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}" + )