From f85e3a550c188a73d34197fbca57433b0aa3ce87 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Mon, 19 Oct 2026 10:41:18 +0100 Subject: [PATCH] Fix zero-step NaN and CPU sync in EDM, Cosine, and Inverse DPMSolverMultistepScheduler Propagates two fixes from DPMSolverMultistepScheduler (scheduling_dpmsolver_multistep.py) to its sibling multistep schedulers: 1. Guard lower_order_final when the last two sigmas are equal (e.g. final_sigmas_type='sigma_min' or algorithm_type in ['sde-dpmsolver++', 'sde-dpmsolver']). Without (len(self.sigmas) > 1 and self.sigmas[-1] == self.sigmas[-2]), taking second- or third-order multistep updates on a zero-width final step divides by h = lambda_t - lambda_s = 0 and corrupts the sample with NaNs. For CosineDPMSolverMultistepScheduler, also guard self.noise_sampler on a zero-width step so BrownianTreeNoiseSampler does not divide by zero (0.0 * NaN = NaN). Unskips the long-standing TestEDMDPMSolverMultistepScheduler.test_solver_order_and_type test. 2. Slice sigmas[-1:] instead of extracting sigmas[-1].item() when final_sigmas_type='sigma_min' in EDMDPMSolverMultistepScheduler and CosineDPMSolverMultistepScheduler, avoiding unnecessary tensor-to-Python scalar extraction during set_timesteps. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera --- .../scheduling_cosine_dpmsolver_multistep.py | 16 +++--- .../scheduling_dpmsolver_multistep_inverse.py | 4 +- .../scheduling_edm_dpmsolver_multistep.py | 7 +-- .../test_scheduler_dpm_multi_inverse.py | 20 +++++++ .../test_scheduler_edm_dpmsolver_multistep.py | 53 ++++++++++++++++++- 5 files changed, 88 insertions(+), 12 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py index 2ba5c377744a..a7e868daa2d7 100644 --- a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py @@ -280,15 +280,15 @@ def set_timesteps(self, num_inference_steps: int | None = None, device: str | to self.timesteps = self.precondition_noise(sigmas) if self.config.final_sigmas_type == "sigma_min": - sigma_last = self.config.sigma_min + sigma_last = sigmas[-1:] elif self.config.final_sigmas_type == "zero": - sigma_last = 0 + sigma_last = torch.zeros(1, dtype=torch.float32, device=device) else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) - self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)]) + self.sigmas = torch.cat([sigmas, sigma_last]) self.model_outputs = [ None, @@ -643,6 +643,7 @@ def step( self.config.euler_at_final or (self.config.lower_order_final and len(self.timesteps) < 15) or self.config.final_sigmas_type == "zero" + or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]) ) lower_order_second = ( (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 @@ -665,9 +666,12 @@ def step( sigma_max=self.config.sigma_max, seed=seed, ) - noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to( - model_output.device - ) + if self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]: + noise = torch.zeros_like(model_output) + else: + noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to( + model_output.device + ) if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py index b236e9975fed..a86e4dd69ee7 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py @@ -1058,7 +1058,9 @@ def step( # Improve numerical stability for small number of steps lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( - self.config.euler_at_final or (self.config.lower_order_final and len(self.timesteps) < 15) + self.config.euler_at_final + or (self.config.lower_order_final and len(self.timesteps) < 15) + or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]) ) lower_order_second = ( (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 diff --git a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py index 20ad17cfdd96..65331f4a9232 100644 --- a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py @@ -303,15 +303,15 @@ def set_timesteps( self.timesteps = self.precondition_noise(sigmas) if self.config.final_sigmas_type == "sigma_min": - sigma_last = self.config.sigma_min + sigma_last = sigmas[-1:] elif self.config.final_sigmas_type == "zero": - sigma_last = 0 + sigma_last = torch.zeros(1, dtype=torch.float32, device=device) else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) - self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)]) + self.sigmas = torch.cat([sigmas, sigma_last]) self.model_outputs = [ None, @@ -766,6 +766,7 @@ def step( self.config.euler_at_final or (self.config.lower_order_final and len(self.timesteps) < 15) or self.config.final_sigmas_type == "zero" + or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]) ) lower_order_second = ( (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 diff --git a/tests/schedulers/test_scheduler_dpm_multi_inverse.py b/tests/schedulers/test_scheduler_dpm_multi_inverse.py index 0eced957190c..4c7503f064a1 100644 --- a/tests/schedulers/test_scheduler_dpm_multi_inverse.py +++ b/tests/schedulers/test_scheduler_dpm_multi_inverse.py @@ -271,3 +271,23 @@ def test_beta_sigmas(self): def test_exponential_sigmas(self): self.check_over_configs(use_exponential_sigmas=True) + + def test_zero_width_final_step(self): + scheduler_class = self.scheduler_classes[0] + model = self.dummy_model() + for kwargs in [ + {"solver_order": 3, "use_karras_sigmas": True}, + {"solver_order": 3, "use_exponential_sigmas": True}, + {"solver_order": 3, "use_beta_sigmas": True}, + {"solver_order": 3, "timestep_spacing": "trailing"}, + {"solver_order": 2, "solver_type": "heun", "use_karras_sigmas": True}, + ]: + scheduler_config = self.get_scheduler_config(**kwargs) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(20) + sample = self.dummy_sample_deter + generator = torch.manual_seed(0) + for t in scheduler.timesteps: + residual = model(sample, t) + sample = scheduler.step(residual, t, sample, generator=generator).prev_sample + assert not torch.isnan(sample).any() diff --git a/tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py b/tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py index 8525ce61c40d..c9706c2523a9 100644 --- a/tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py +++ b/tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py @@ -165,8 +165,6 @@ def test_prediction_type(self): for prediction_type in ["epsilon", "v_prediction"]: self.check_over_configs(prediction_type=prediction_type) - # TODO (patil-suraj): Fix this test - @unittest.skip("Skip for now, as it failing currently but works with the actual model") def test_solver_order_and_type(self): for algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: for solver_type in ["midpoint", "heun"]: @@ -260,3 +258,54 @@ def test_duplicated_timesteps(self, **config): @unittest.skip("Test not supported.") def test_trained_betas(self): pass + + def test_zero_width_final_step(self): + scheduler_class = self.scheduler_classes[0] + model = self.dummy_model() + for kwargs in [ + {"solver_order": 3, "sigma_schedule": "karras", "final_sigmas_type": "sigma_min"}, + {"solver_order": 3, "sigma_schedule": "exponential", "final_sigmas_type": "sigma_min"}, + {"solver_order": 2, "solver_type": "heun", "sigma_schedule": "karras", "final_sigmas_type": "sigma_min"}, + { + "solver_order": 2, + "solver_type": "heun", + "algorithm_type": "sde-dpmsolver++", + "sigma_min": 0.01, + "final_sigmas_type": "sigma_min", + }, + { + "solver_order": 1, + "algorithm_type": "sde-dpmsolver++", + "sigma_min": 0.01, + "final_sigmas_type": "sigma_min", + }, + ]: + scheduler_config = self.get_scheduler_config(**kwargs) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(20) + sample = self.dummy_sample_deter + generator = torch.manual_seed(0) + for t in scheduler.timesteps: + residual = model(sample, t) + sample = scheduler.step(residual, t, sample, generator=generator).prev_sample + assert not torch.isnan(sample).any() + + from diffusers import CosineDPMSolverMultistepScheduler + + for solver_order in [1, 2]: + for solver_type in ["midpoint", "heun"]: + cosine_scheduler = CosineDPMSolverMultistepScheduler( + solver_order=solver_order, + solver_type=solver_type, + final_sigmas_type="sigma_min", + ) + cosine_scheduler.set_timesteps(10) + sample = self.dummy_sample_deter + # BrownianTreeNoiseSampler divides by sqrt(|t1 - t0|) and yields NaN when sigma == sigma_next + cosine_scheduler.noise_sampler = lambda s, s_next: ( + torch.full_like(sample, float("nan")) if s == s_next else torch.zeros_like(sample) + ) + for t in cosine_scheduler.timesteps: + residual = model(sample, t) + sample = cosine_scheduler.step(residual, t, sample).prev_sample + assert not torch.isnan(sample).any()