From 67999193ff519116e5bfd1e48c4ca40fcbaf2c4a Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Fri, 16 Sep 2022 17:47:25 +0200 Subject: [PATCH] Revert "adding more typehints to DDIM scheduler (#456)" This reverts commit a0558b11465d446ce0b5d0e6e9c632b33ea7a542. --- src/diffusers/schedulers/scheduling_ddim.py | 30 ++++++++++----------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index c1b3cb410c05..85ba67523172 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 List, Optional, Tuple, Union +from typing import 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: int, max_beta: float = 0.999) -> np.ndarray: +def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=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: 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) -> float: + def alpha_bar(time_step): return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2 - 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)) + 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)) 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.ndarray = np.asarray(trained_betas) + self.betas = 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: np.ndarray = 1.0 - self.betas - self.alphas_cumprod: np.ndarray = np.cumprod(self.alphas, axis=0) + self.alphas = 1.0 - self.betas + self.alphas_cumprod = 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: int = 0 - self.timesteps: np.ndarray = np.arange(0, num_train_timesteps)[::-1].copy() + self.num_inference_steps = None + self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy() - self.tensor_format: str = tensor_format + self.tensor_format = 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) -> None: + def set_timesteps(self, num_inference_steps: int, offset: int = 0): """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.