[UniPCMultistepScheduler] Fix final-step NaN/inf at sigma=0 and rks dtype mismatch

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -275,6 +275,11 @@ def __init__(
             else:
                 raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
 
+        if not predict_x0 and final_sigmas_type == "zero":
+            raise ValueError(
+                "`final_sigmas_type` zero is not supported for `predict_x0=False`. Please choose `sigma_min` instead."
+            )
+
         self.predict_x0 = predict_x0
         # setable values
         self.num_inference_steps = None
@@ -903,7 +908,7 @@ def multistep_uni_p_bh_update(
             rks.append(rk)
             D1s.append((mi - m0) / rk)
 
-        rks.append(torch.ones((), device=device))
+        rks.append(torch.ones_like(h))
         rks = torch.stack(rks)
 
         R = []
@@ -945,16 +950,16 @@ def multistep_uni_p_bh_update(
             x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
             if D1s is not None:
                 pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
+                x_t = x_t_ - alpha_t * B_h * pred_res
             else:
-                pred_res = 0
-            x_t = x_t_ - alpha_t * B_h * pred_res
+                x_t = x_t_
         else:
             x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
             if D1s is not None:
                 pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
+                x_t = x_t_ - sigma_t * B_h * pred_res
             else:
-                pred_res = 0
-            x_t = x_t_ - sigma_t * B_h * pred_res
+                x_t = x_t_
 
         x_t = x_t.to(x.dtype)
         return x_t
@@ -1038,7 +1043,7 @@ def multistep_uni_c_bh_update(
             rks.append(rk)
             D1s.append((mi - m0) / rk)
 
-        rks.append(torch.ones((), device=device))
+        rks.append(torch.ones_like(h))
         rks = torch.stack(rks)
 
         R = []
@@ -1205,7 +1210,9 @@ def step(
         self.model_outputs[-1] = model_output_convert
         self.timestep_list[-1] = timestep
 
-        if self.config.lower_order_final:
+        if self.config.final_sigmas_type == "zero" and self.step_index == len(self.timesteps) - 1:
+            this_order = 1
+        elif self.config.lower_order_final:
             this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
         else:
             this_order = self.config.solver_order

tests/schedulers/test_scheduler_unipc.py

@@ -284,6 +284,61 @@ def test_full_loop_with_noise(self):
         assert abs(result_sum.item() - 315.5757) < 1e-2, f" expected result sum 315.5757, but get {result_sum}"
         assert abs(result_mean.item() - 0.4109) < 1e-3, f" expected result mean 0.4109, but get {result_mean}"
 
+    def test_last_step_to_sigma_zero_is_the_data_prediction(self):
+        scheduler_class = self.scheduler_classes[0]
+        model = self.dummy_model()
+
+        for solver_type in ["bh1", "bh2"]:
+            for order in [1, 2, 3]:
+                for lower_order_final in [True, False]:
+                    scheduler_config = self.get_scheduler_config(
+                        solver_type=solver_type,
+                        solver_order=order,
+                        lower_order_final=lower_order_final,
+                        final_sigmas_type="zero",
+                    )
+                    scheduler = scheduler_class(**scheduler_config)
+                    scheduler.set_timesteps(10)
+
+                    sample = self.dummy_sample_deter.clone()
+                    for t in scheduler.timesteps[:-1]:
+                        sample = scheduler.step(model(sample, t), t, sample).prev_sample
+
+                    t_last = scheduler.timesteps[-1]
+                    last_residual = model(sample, t_last)
+                    expected = scheduler.convert_model_output(last_residual, sample=sample)
+                    prev_sample = scheduler.step(last_residual, t_last, sample).prev_sample
+
+                    assert torch.isfinite(prev_sample).all()
+                    assert torch.allclose(prev_sample, expected, atol=1e-6)
+
+    def test_predict_x0_false_requires_final_sigmas_type_sigma_min(self):
+        scheduler_class = self.scheduler_classes[0]
+        scheduler_config = self.get_scheduler_config(predict_x0=False, final_sigmas_type="zero")
+        with self.assertRaises(ValueError):
+            scheduler_class(**scheduler_config)
+
+    def test_default_dtype_float64(self):
+        scheduler_class = self.scheduler_classes[0]
+        scheduler_config = self.get_scheduler_config(solver_order=2)
+        scheduler = scheduler_class(**scheduler_config)
+        scheduler.set_timesteps(10)
+
+        model = self.dummy_model()
+        sample = self.dummy_sample_deter.clone()
+
+        prev_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(torch.float64)
+        try:
+            for t in scheduler.timesteps:
+                residual = model(sample, t)
+                sample = scheduler.step(residual, t, sample).prev_sample
+        finally:
+            torch.set_default_dtype(prev_dtype)
+
+        assert sample.dtype == torch.float32
+        assert torch.isfinite(sample).all()
+
 
 class UniPCMultistepScheduler1DTest(UniPCMultistepSchedulerTest):
     @property