From 83e3373ac3f32c0cadaa20a0b0d94eb463d749cd Mon Sep 17 00:00:00 2001 From: V Vishnu Anirudh Date: Fri, 9 Sep 2022 21:44:04 +0100 Subject: [PATCH 1/4] adding more typehints --- src/diffusers/schedulers/scheduling_ddim.py | 26 ++++++++++----------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index 894d63bf2df0..baf4e6374a94 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -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): """ 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): 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) @@ -97,20 +97,20 @@ 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) + self.betas: np.ndarray = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2 + self.betas: np.ndarray = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2 elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) + self.betas: np.ndarray = betas_for_alpha_bar(num_train_timesteps) 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: float = 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 From f7c7b532d8507f90a0f0b7942bcd811dde79c420 Mon Sep 17 00:00:00 2001 From: V Vishnu Anirudh Date: Fri, 9 Sep 2022 22:29:29 +0100 Subject: [PATCH 2/4] resolving mypy issues --- src/diffusers/schedulers/scheduling_ddim.py | 24 ++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index baf4e6374a94..dcc740cdd832 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/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 Optional, Tuple, Union, List import numpy as np import torch @@ -25,7 +25,7 @@ from .scheduling_utils import SchedulerMixin, SchedulerOutput -def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta: float=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,10 +43,10 @@ def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta: float=0.999): betas (`np.ndarray`): the betas used by the scheduler to step the model outputs """ - def calculate_alpha_bar(time_step: float): + def calculate_alpha_bar(time_step: float) -> float: return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2 - betas: list[float] = [] + 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 @@ -99,17 +99,17 @@ def __init__( if trained_betas is not None: self.betas: np.ndarray = np.asarray(trained_betas) if beta_schedule == "linear": - self.betas: np.ndarray = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) + self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas: np.ndarray = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2 + self.betas = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2 elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule - self.betas: np.ndarray = betas_for_alpha_bar(num_train_timesteps) + self.betas = betas_for_alpha_bar(num_train_timesteps) else: raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") - self.alphas: float = 1.0 - self.betas + 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 @@ -119,10 +119,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): @@ -135,7 +135,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. From a4db7b365faa3f4e94ef1298b9c124fcfa45bdc0 Mon Sep 17 00:00:00 2001 From: V Vishnu Anirudh Date: Fri, 9 Sep 2022 22:34:32 +0100 Subject: [PATCH 3/4] resolving formatting issue --- src/diffusers/schedulers/scheduling_ddim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index dcc740cdd832..a536e5f679fa 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -25,7 +25,7 @@ from .scheduling_utils import SchedulerMixin, SchedulerOutput -def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta: float=0.999)-> np.ndarray: +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]. From e602dbd242cf8290bad42b2205e9dbe3b56bf7b1 Mon Sep 17 00:00:00 2001 From: V Vishnu Anirudh Date: Sun, 11 Sep 2022 01:23:41 +0100 Subject: [PATCH 4/4] fixing isort issue --- src/diffusers/schedulers/scheduling_ddim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index a536e5f679fa..01f62cc54269 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -16,7 +16,7 @@ # and https://github.com/hojonathanho/diffusion import math -from typing import Optional, Tuple, Union, List +from typing import List, Optional, Tuple, Union import numpy as np import torch