Revert "adding more typehints to DDIM scheduler" (#533)

Revert "adding more typehints to DDIM scheduler (#456)"

This reverts commit a0558b1146.
This commit is contained in:
Patrick von Platen 2022-09-16 17:48:02 +02:00 committed by GitHub
parent a0558b1146
commit 88972172d8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 15 deletions

View File

@ -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 ..configuration_utils import ConfigMixin, register_to_config
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 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin):
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 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin):
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 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin):
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 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin):
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.