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:
parent
a0558b1146
commit
88972172d8
|
@ -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.
|
||||
|
||||
|
|
Loading…
Reference in New Issue