Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 15 additions & 15 deletions src/diffusers/schedulers/scheduling_ddim.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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].
Expand All @@ -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)


Expand Down Expand Up @@ -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":
Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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.

Expand Down