[Type hint] PNDM schedulers (#335)

* [Type hint] PNDM Schedulers

* ran make style

* updated timesteps type hint

* apply suggestions from code review

* ran make style

* removed unused import

Author: daspartho

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)
@@ -120,7 +120,7 @@ def __init__(
         self.tensor_format = tensor_format
         self.set_format(tensor_format=tensor_format)
 
-    def set_timesteps(self, num_inference_steps, offset=0):
+    def set_timesteps(self, num_inference_steps: int, offset: int = 0) -> torch.FloatTensor:
         self.num_inference_steps = num_inference_steps
         self._timesteps = list(
             range(0, self.config.num_train_timesteps, self.config.num_train_timesteps // num_inference_steps)
@@ -287,7 +287,13 @@ def _get_prev_sample(self, sample, timestep, timestep_prev, model_output):
 
         return prev_sample
 
-    def add_noise(self, original_samples, noise, timesteps):
+    def add_noise(
+        self,
+        original_samples: Union[torch.FloatTensor, np.ndarray],
+        noise: Union[torch.FloatTensor, np.ndarray],
+        timesteps: Union[torch.IntTensor, np.ndarray],
+    ) -> torch.Tensor:
+
         sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
         sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
         sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5