diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index edcfa4576e53..7bfcf5c2b94a 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -1240,6 +1240,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/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..5d6325646867 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -1207,6 +1207,10 @@ def step( if self.config.lower_order_final: this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) + elif self.step_index == len(self.timesteps) - 1 and bool( + self.sigmas[self.step_index + 1] == self.sigmas[self.step_index] + ): + this_order = 1 else: this_order = self.config.solver_order diff --git a/tests/schedulers/test_scheduler_dpm_multi.py b/tests/schedulers/test_scheduler_dpm_multi.py index 28c354709dc9..ae0e68d7d795 100644 --- a/tests/schedulers/test_scheduler_dpm_multi.py +++ b/tests/schedulers/test_scheduler_dpm_multi.py @@ -366,3 +366,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, "final_sigmas_type": "sigma_min"}, + {"solver_order": 3, "use_exponential_sigmas": True, "final_sigmas_type": "sigma_min"}, + {"solver_order": 3, "use_beta_sigmas": True, "final_sigmas_type": "sigma_min"}, + {"solver_order": 3, "use_lu_lambdas": True, "final_sigmas_type": "sigma_min"}, + {"solver_order": 2, "solver_type": "heun", "use_karras_sigmas": True, "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() diff --git a/tests/schedulers/test_scheduler_unipc.py b/tests/schedulers/test_scheduler_unipc.py index ac7e1d3f88b4..c03822bec91a 100644 --- a/tests/schedulers/test_scheduler_unipc.py +++ b/tests/schedulers/test_scheduler_unipc.py @@ -428,3 +428,25 @@ def test_flow_and_karras_sigmas_values(self): expected_timesteps = expected_timesteps[0:-1] self.assertTrue(torch.allclose(scheduler.sigmas, expected_sigmas)) self.assertTrue(torch.all(expected_timesteps == scheduler.timesteps)) + + def test_zero_width_final_step(self): + scheduler_class = self.scheduler_classes[0] + model = self.dummy_model() + for kwargs in [ + {"use_karras_sigmas": True}, + {"use_exponential_sigmas": True}, + {"use_beta_sigmas": True}, + ]: + scheduler_config = self.get_scheduler_config( + solver_order=3, + lower_order_final=False, + final_sigmas_type="sigma_min", + **kwargs, + ) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(20) + sample = self.dummy_sample_deter + for t in scheduler.timesteps: + residual = model(sample, t) + sample = scheduler.step(residual, t, sample).prev_sample + assert not torch.isnan(sample).any()