|
from abc import abstractmethod |
|
from functools import partial |
|
|
|
import numpy as np |
|
import torch |
|
|
|
from ...modules.diffusionmodules.util import make_beta_schedule |
|
from ...util import append_zero |
|
|
|
|
|
def generate_roughly_equally_spaced_steps( |
|
num_substeps: int, max_step: int |
|
) -> np.ndarray: |
|
return np.linspace(max_step - 1, 0, num_substeps, endpoint=False).astype(int)[::-1] |
|
|
|
|
|
class Discretization: |
|
def __call__(self, n, do_append_zero=True, device="cpu", flip=False): |
|
sigmas = self.get_sigmas(n, device=device) |
|
sigmas = append_zero(sigmas) if do_append_zero else sigmas |
|
return sigmas if not flip else torch.flip(sigmas, (0,)) |
|
|
|
@abstractmethod |
|
def get_sigmas(self, n, device): |
|
pass |
|
|
|
|
|
class EDMDiscretization(Discretization): |
|
def __init__(self, sigma_min=0.002, sigma_max=80.0, rho=7.0): |
|
self.sigma_min = sigma_min |
|
self.sigma_max = sigma_max |
|
self.rho = rho |
|
|
|
def get_sigmas(self, n, device="cpu"): |
|
ramp = torch.linspace(0, 1, n, device=device) |
|
min_inv_rho = self.sigma_min ** (1 / self.rho) |
|
max_inv_rho = self.sigma_max ** (1 / self.rho) |
|
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** self.rho |
|
return sigmas |
|
|
|
|
|
class LegacyDDPMDiscretization(Discretization): |
|
def __init__( |
|
self, |
|
linear_start=0.00085, |
|
linear_end=0.0120, |
|
num_timesteps=1000, |
|
): |
|
super().__init__() |
|
self.num_timesteps = num_timesteps |
|
betas = make_beta_schedule( |
|
"linear", num_timesteps, linear_start=linear_start, linear_end=linear_end |
|
) |
|
alphas = 1.0 - betas |
|
self.alphas_cumprod = np.cumprod(alphas, axis=0) |
|
self.to_torch = partial(torch.tensor, dtype=torch.float32) |
|
|
|
def get_sigmas(self, n, device="cpu"): |
|
if n < self.num_timesteps: |
|
timesteps = generate_roughly_equally_spaced_steps(n, self.num_timesteps) |
|
alphas_cumprod = self.alphas_cumprod[timesteps] |
|
elif n == self.num_timesteps: |
|
alphas_cumprod = self.alphas_cumprod |
|
else: |
|
raise ValueError |
|
|
|
to_torch = partial(torch.tensor, dtype=torch.float32, device=device) |
|
sigmas = to_torch((1 - alphas_cumprod) / alphas_cumprod) ** 0.5 |
|
return torch.flip(sigmas, (0,)) |
|
|
