adding more typehints to DDIM scheduler (open-mmlab#456)

* adding more typehints

* resolving mypy issues

* resolving formatting issue

* fixing isort issue

Co-authored-by: V Vishnu Anirudh <[email protected]>
Co-authored-by: V Vishnu Anirudh <[email protected]>

src/diffusers/schedulers/scheduling_ddim.py

@@ -16,7 +16,7 @@
 # and https://github.com/hojonathanho/diffusion
 
 import math
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -25,7 +25,7 @@
 from .scheduling_utils import SchedulerMixin, SchedulerOutput
 
 
-def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
+def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta: float = 0.999) -> np.ndarray:
     """
     Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
     (1-beta) over time from t = [0,1].
@@ -43,14 +43,14 @@ def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
         betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
     """
 
-    def alpha_bar(time_step):
+    def calculate_alpha_bar(time_step: float) -> float:
         return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2
 
-    betas = []
-    for i in range(num_diffusion_timesteps):
-        t1 = i / num_diffusion_timesteps
-        t2 = (i + 1) / num_diffusion_timesteps
-        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    betas: List[float] = []
+    for diffusion_timestep in range(num_diffusion_timesteps):
+        lower_timestep = diffusion_timestep / num_diffusion_timesteps
+        upper_timestep = (diffusion_timestep + 1) / num_diffusion_timesteps
+        betas.append(min(1 - calculate_alpha_bar(upper_timestep) / calculate_alpha_bar(lower_timestep), max_beta))
     return np.array(betas, dtype=np.float32)
 
 
@@ -96,7 +96,7 @@ def __init__(
         tensor_format: str = "pt",
     ):
         if trained_betas is not None:
-            self.betas = np.asarray(trained_betas)
+            self.betas: np.ndarray = np.asarray(trained_betas)
         if beta_schedule == "linear":
             self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
         elif beta_schedule == "scaled_linear":
@@ -108,8 +108,8 @@ def __init__(
         else:
             raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
 
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
+        self.alphas: np.ndarray = 1.0 - self.betas
+        self.alphas_cumprod: np.ndarray = np.cumprod(self.alphas, axis=0)
 
         # At every step in ddim, we are looking into the previous alphas_cumprod
         # For the final step, there is no previous alphas_cumprod because we are already at 0
@@ -118,10 +118,10 @@ def __init__(
         self.final_alpha_cumprod = np.array(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
 
         # setable values
-        self.num_inference_steps = None
-        self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy()
+        self.num_inference_steps: int = 0
+        self.timesteps: np.ndarray = np.arange(0, num_train_timesteps)[::-1].copy()
 
-        self.tensor_format = tensor_format
+        self.tensor_format: str = tensor_format
         self.set_format(tensor_format=tensor_format)
 
     def _get_variance(self, timestep, prev_timestep):
@@ -134,7 +134,7 @@ def _get_variance(self, timestep, prev_timestep):
 
         return variance
 
-    def set_timesteps(self, num_inference_steps: int, offset: int = 0):
+    def set_timesteps(self, num_inference_steps: int, offset: int = 0) -> None:
         """
         Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.