Fix beta and exponential sigmas + add tests (#9954)

* Fix beta and exponential sigmas + add tests

---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

tests/schedulers/test_scheduler_kdpm2_discrete.py

@@ -164,3 +164,9 @@ class KDPM2DiscreteSchedulerTest(SchedulerCommonTest):
 
         assert abs(result_sum.item() - 70408.4062) < 1e-2, f" expected result sum 70408.4062, but get {result_sum}"
         assert abs(result_mean.item() - 91.6776) < 1e-3, f" expected result mean 91.6776, but get {result_mean}"
+
+    def test_beta_sigmas(self):
+        self.check_over_configs(use_beta_sigmas=True)
+
+    def test_exponential_sigmas(self):
+        self.check_over_configs(use_exponential_sigmas=True)

tests/schedulers/test_scheduler_lms.py

@@ -168,3 +168,9 @@ class LMSDiscreteSchedulerTest(SchedulerCommonTest):
 
         assert abs(result_sum.item() - 27663.6895) < 1e-2
         assert abs(result_mean.item() - 36.0204) < 1e-3
+
+    def test_beta_sigmas(self):
+        self.check_over_configs(use_beta_sigmas=True)
+
+    def test_exponential_sigmas(self):
+        self.check_over_configs(use_exponential_sigmas=True)

tests/schedulers/test_scheduler_sasolver.py

@@ -200,3 +200,9 @@ class SASolverSchedulerTest(SchedulerCommonTest):
             assert abs(result_mean.item() - 1.0901763439178467) < 1e-2
         else:
             print("None")
+
+    def test_beta_sigmas(self):
+        self.check_over_configs(use_beta_sigmas=True)
+
+    def test_exponential_sigmas(self):
+        self.check_over_configs(use_exponential_sigmas=True)

tests/schedulers/test_scheduler_unipc.py

@@ -393,3 +393,9 @@ class UniPCMultistepScheduler1DTest(UniPCMultistepSchedulerTest):
 
         assert abs(result_sum.item() - 39.0870) < 1e-2, f" expected result sum 39.0870, but get {result_sum}"
         assert abs(result_mean.item() - 0.4072) < 1e-3, f" expected result mean 0.4072, but get {result_mean}"
+
+    def test_beta_sigmas(self):
+        self.check_over_configs(use_beta_sigmas=True)
+
+    def test_exponential_sigmas(self):
+        self.check_over_configs(use_exponential_sigmas=True)