Fix NaN corruption on final step in DPMSolverMultistep and UniPCMultistep with final_sigmas_type="sigma_min"

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

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
 

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()

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()