diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 8ad6abd90668..ec6a69f80aad 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/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 diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py index 706190e291a4..5accffafc175 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +++ b/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 diff --git a/tests/schedulers/test_scheduler_euler_ancestral.py b/tests/schedulers/test_scheduler_euler_ancestral.py index c4fe61bfc387..48e41035ae1b 100644 --- a/tests/schedulers/test_scheduler_euler_ancestral.py +++ b/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'" diff --git a/tests/schedulers/test_scheduler_kdpm2_ancestral.py b/tests/schedulers/test_scheduler_kdpm2_ancestral.py index 135534db4536..92410ec66f09 100644 --- a/tests/schedulers/test_scheduler_kdpm2_ancestral.py +++ b/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'" + )