[Type hint] PNDM schedulers (#335)

daspartho · natolambert · commit 6756d18d8e1b · 2022-09-07T11:42:28.000-06:00
* [Type hint] PNDM Schedulers

* ran make style

* updated timesteps type hint

* apply suggestions from code review

* ran make style

* removed unused import
diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py
@@ -72,13 +72,13 @@ class PNDMScheduler(SchedulerMixin, ConfigMixin):
     @register_to_config
     def __init__(
         self,
-        num_train_timesteps=1000,
-        beta_start=0.0001,
-        beta_end=0.02,
-        beta_schedule="linear",
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
         trained_betas=None,
-        tensor_format="pt",
-        skip_prk_steps=False,
+        tensor_format: str = "pt",
+        skip_prk_steps: bool = False,
     ):
         if trained_betas is not None:
             self.betas = np.asarray(trained_betas)
