Fix NaN in KDPM2DiscreteScheduler.sigmas_interpol (all-NaN output on MPS)

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(

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}"
+            )