refactor DPMSolverMultistepScheduler using sigmas (#4986)

---------
Co-authored-by: yiyixuxu <yixu310@gmail,com>
Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

tests/schedulers/test_scheduler_deis.py

@@ -51,6 +51,7 @@ class DEISMultistepSchedulerTest(SchedulerCommonTest):
 
             output, new_output = sample, sample
             for t in range(time_step, time_step + scheduler.config.solver_order + 1):
+                t = scheduler.timesteps[t]
                 output = scheduler.step(residual, t, output, **kwargs).prev_sample
                 new_output = new_scheduler.step(residual, t, new_output, **kwargs).prev_sample
 

tests/schedulers/test_scheduler_dpm_multi.py

@@ -59,6 +59,7 @@ class DPMSolverMultistepSchedulerTest(SchedulerCommonTest):
 
             output, new_output = sample, sample
             for t in range(time_step, time_step + scheduler.config.solver_order + 1):
+                t = new_scheduler.timesteps[t]
                 output = scheduler.step(residual, t, output, **kwargs).prev_sample
                 new_output = new_scheduler.step(residual, t, new_output, **kwargs).prev_sample
 
@@ -91,6 +92,7 @@ class DPMSolverMultistepSchedulerTest(SchedulerCommonTest):
                 # copy over dummy past residual (must be after setting timesteps)
                 new_scheduler.model_outputs = dummy_past_residuals[: new_scheduler.config.solver_order]
 
+            time_step = new_scheduler.timesteps[time_step]
             output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample
             new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample
 
@@ -264,10 +266,10 @@ class DPMSolverMultistepSchedulerTest(SchedulerCommonTest):
 
         assert sample.dtype == torch.float16
 
-    def test_unique_timesteps(self, **config):
+    def test_duplicated_timesteps(self, **config):
         for scheduler_class in self.scheduler_classes:
             scheduler_config = self.get_scheduler_config(**config)
             scheduler = scheduler_class(**scheduler_config)
 
             scheduler.set_timesteps(scheduler.config.num_train_timesteps)
-            assert len(scheduler.timesteps.unique()) == scheduler.num_inference_steps
+            assert len(scheduler.timesteps) == scheduler.num_inference_steps

tests/schedulers/test_scheduler_dpm_multi_inverse.py

@@ -54,6 +54,7 @@ class DPMSolverMultistepSchedulerTest(SchedulerCommonTest):
 
             output, new_output = sample, sample
             for t in range(time_step, time_step + scheduler.config.solver_order + 1):
+                t = scheduler.timesteps[t]
                 output = scheduler.step(residual, t, output, **kwargs).prev_sample
                 new_output = new_scheduler.step(residual, t, new_output, **kwargs).prev_sample
 
@@ -222,7 +223,7 @@ class DPMSolverMultistepSchedulerTest(SchedulerCommonTest):
         sample = self.full_loop(prediction_type="v_prediction", use_karras_sigmas=True)
         result_mean = torch.mean(torch.abs(sample))
 
-        assert abs(result_mean.item() - 1.7833) < 1e-3
+        assert abs(result_mean.item() - 1.7833) < 2e-3
 
     def test_switch(self):
         # make sure that iterating over schedulers with same config names gives same results

tests/schedulers/test_scheduler_dpm_single.py

@@ -58,6 +58,7 @@ class DPMSolverSinglestepSchedulerTest(SchedulerCommonTest):
 
             output, new_output = sample, sample
             for t in range(time_step, time_step + scheduler.config.solver_order + 1):
+                t = scheduler.timesteps[t]
                 output = scheduler.step(residual, t, output, **kwargs).prev_sample
                 new_output = new_scheduler.step(residual, t, new_output, **kwargs).prev_sample
 
@@ -248,3 +249,33 @@ class DPMSolverSinglestepSchedulerTest(SchedulerCommonTest):
             sample = scheduler.step(residual, t, sample).prev_sample
 
         assert sample.dtype == torch.float16
+
+    def test_step_shape(self):
+        kwargs = dict(self.forward_default_kwargs)
+
+        num_inference_steps = kwargs.pop("num_inference_steps", None)
+
+        for scheduler_class in self.scheduler_classes:
+            scheduler_config = self.get_scheduler_config()
+            scheduler = scheduler_class(**scheduler_config)
+
+            sample = self.dummy_sample
+            residual = 0.1 * sample
+
+            if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"):
+                scheduler.set_timesteps(num_inference_steps)
+            elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"):
+                kwargs["num_inference_steps"] = num_inference_steps
+
+            # copy over dummy past residuals (must be done after set_timesteps)
+            dummy_past_residuals = [residual + 0.2, residual + 0.15, residual + 0.10]
+            scheduler.model_outputs = dummy_past_residuals[: scheduler.config.solver_order]
+
+            time_step_0 = scheduler.timesteps[0]
+            time_step_1 = scheduler.timesteps[1]
+
+            output_0 = scheduler.step(residual, time_step_0, sample, **kwargs).prev_sample
+            output_1 = scheduler.step(residual, time_step_1, sample, **kwargs).prev_sample
+
+            self.assertEqual(output_0.shape, sample.shape)
+            self.assertEqual(output_0.shape, output_1.shape)

tests/schedulers/test_scheduler_unipc.py

@@ -52,6 +52,7 @@ class UniPCMultistepSchedulerTest(SchedulerCommonTest):
 
             output, new_output = sample, sample
             for t in range(time_step, time_step + scheduler.config.solver_order + 1):
+                t = scheduler.timesteps[t]
                 output = scheduler.step(residual, t, output, **kwargs).prev_sample
                 new_output = new_scheduler.step(residual, t, new_output, **kwargs).prev_sample
 
@@ -241,11 +242,3 @@ class UniPCMultistepSchedulerTest(SchedulerCommonTest):
             sample = scheduler.step(residual, t, sample).prev_sample
 
         assert sample.dtype == torch.float16
-
-    def test_unique_timesteps(self, **config):
-        for scheduler_class in self.scheduler_classes:
-            scheduler_config = self.get_scheduler_config(**config)
-            scheduler = scheduler_class(**scheduler_config)
-
-            scheduler.set_timesteps(scheduler.config.num_train_timesteps)
-            assert len(scheduler.timesteps.unique()) == scheduler.num_inference_steps