fix(schedulers): prevent NaN in ancestral schedulers with squaredcos_cap_v2

src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py

@@ -445,7 +445,12 @@ def step(
         sigma_from = self.sigmas[self.step_index]
         sigma_to = self.sigmas[self.step_index + 1]
         sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
-        sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
+        # Clamp the radicand to 0. sigma_up is mathematically <= sigma_to, but with
+        # beta_schedule="squaredcos_cap_v2" the sigma dynamic range is so large that in
+        # float32 sigma_up can round to a value fractionally larger than sigma_to, making
+        # sigma_to**2 - sigma_up**2 negative and producing NaN. sigma_up == sigma_to is the
+        # correct limit, so sigma_down == 0 there.
+        sigma_down = ((sigma_to**2 - sigma_up**2).clamp(min=0)) ** 0.5
 
         # 2. Convert to an ODE derivative
         derivative = (sample - pred_original_sample) / sigma

src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py

@@ -310,12 +310,28 @@ def set_timesteps(
         # compute up and down sigmas
         sigmas_next = sigmas.roll(-1)
         sigmas_next[-1] = 0.0
-        sigmas_up = (sigmas_next**2 * (sigmas**2 - sigmas_next**2) / sigmas**2) ** 0.5
-        sigmas_down = (sigmas_next**2 - sigmas_up**2) ** 0.5
-        sigmas_down[-1] = 0.0
+        # Clamp the radicand to 0. sigmas_up is mathematically <= sigmas_next, but under
+        # beta_schedule="squaredcos_cap_v2" the sigma dynamic range is so large that in
+        # float32 sigmas_up can round to a value fractionally larger than sigmas_next,
+        # making sigmas_next**2 - sigmas_up**2 negative and producing NaN.
+        sigmas_up_sq = sigmas_next**2 * (sigmas**2 - sigmas_next**2) / sigmas**2
+        sigmas_up = sigmas_up_sq.clamp(min=0) ** 0.5
+        # Compute sigmas_down algebraically as sqrt(sigmas_next**4 / sigmas**2). The textbook
+        # form sqrt(sigmas_next**2 - sigmas_up**2) suffers catastrophic cancellation in the
+        # same squaredcos_cap_v2 regime (sigmas_up ~= sigmas_next), which is what drives
+        # sigmas_down to 0 and then NaN into sigmas_interpol via log(0). The algebraic form
+        # is mathematically identical and numerically stable.
+        sigmas_down = sigmas_next**2 / sigmas.clamp(min=1e-30)  # sqrt(sigmas_next**4/sigmas**2)
+        sigmas_down = torch.where(sigmas_next > 0, sigmas_down, torch.zeros_like(sigmas_down))
 
         # compute interpolated sigmas
-        sigmas_interpol = sigmas.log().lerp(sigmas_down.log(), 0.5).exp()
+        # Midpoint between sigmas and sigmas_down in log space. Written as the equivalent
+        # geometric mean (sqrt(sigmas * sigmas_down)) rather than
+        # sigmas.log().lerp(sigmas_down.log(), 0.5).exp() because the latter produces NaN
+        # wherever sigmas_down == 0 (torch.lerp with a -inf end term), which happens at the
+        # tail and — under squaredcos_cap_v2 — at the head where the huge sigma range forces
+        # sigmas_down[0] to 0.
+        sigmas_interpol = (sigmas.clamp(min=0) * sigmas_down) ** 0.5
         sigmas_interpol[-2:] = 0.0
 
         # set sigmas

tests/schedulers/test_scheduler_euler_ancestral.py

@@ -154,3 +154,25 @@ def test_full_loop_with_noise(self):
 
         assert abs(result_sum.item() - 56163.0508) < 1e-2, f" expected result sum 56163.0508, but get {result_sum}"
         assert abs(result_mean.item() - 73.1290) < 1e-3, f" expected result mean  73.1290, but get {result_mean}"
+
+    def test_full_loop_squaredcos_cap_v2_no_nan(self):
+        # Regression test for #14213: with beta_schedule="squaredcos_cap_v2" the
+        # sigma dynamic range is large enough that float32 rounding makes
+        # sigma_up fractionally exceed sigma_to, driving sigma_down's radicand
+        # negative and producing all-NaN output at low step counts.
+        scheduler_class = self.scheduler_classes[0]
+        scheduler_config = self.get_scheduler_config(beta_schedule="squaredcos_cap_v2")
+        scheduler = scheduler_class(**scheduler_config)
+
+        scheduler.set_timesteps(4)
+
+        sample = self.dummy_sample_deter * scheduler.init_noise_sigma
+        model = self.dummy_model()
+
+        for t in scheduler.timesteps:
+            sample = scheduler.scale_model_input(sample, t)
+            model_output = model(sample, t)
+            output = scheduler.step(model_output, t, sample)
+            sample = output.prev_sample
+
+        assert torch.isfinite(sample).all(), "step() produced non-finite output with beta_schedule='squaredcos_cap_v2'"

tests/schedulers/test_scheduler_kdpm2_ancestral.py

@@ -162,3 +162,20 @@ def test_beta_sigmas(self):
 
     def test_exponential_sigmas(self):
         self.check_over_configs(use_exponential_sigmas=True)
+
+    def test_set_timesteps_squaredcos_cap_v2_no_nan(self):
+        # Regression test for #14213: with beta_schedule="squaredcos_cap_v2" the
+        # huge sigma dynamic range made sigmas_down round negative in float32,
+        # which drove log(0) NaN into sigmas_interpol during set_timesteps().
+        scheduler_class = self.scheduler_classes[0]
+        scheduler_config = self.get_scheduler_config(beta_schedule="squaredcos_cap_v2")
+        scheduler = scheduler_class(**scheduler_config)
+
+        scheduler.set_timesteps(4)
+
+        assert torch.isfinite(scheduler.timesteps.float()).all(), (
+            "set_timesteps() produced non-finite timesteps with beta_schedule='squaredcos_cap_v2'"
+        )
+        assert torch.isfinite(scheduler.sigmas).all(), (
+            "set_timesteps() produced non-finite sigmas with beta_schedule='squaredcos_cap_v2'"
+        )