87 lines
3.6 KiB
Python
87 lines
3.6 KiB
Python
import math
|
|
|
|
import torch
|
|
import numpy as np
|
|
|
|
|
|
def make_beta_schedule(device, schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
|
if schedule == "linear":
|
|
betas = (
|
|
torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
|
|
)
|
|
|
|
elif schedule == "cosine":
|
|
timesteps = (torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s).to(device)
|
|
alphas = timesteps / (1 + cosine_s) * np.pi / 2
|
|
alphas = torch.cos(alphas).pow(2).to(device)
|
|
alphas = alphas / alphas[0]
|
|
betas = 1 - alphas[1:] / alphas[:-1]
|
|
betas = np.clip(betas, a_min=0, a_max=0.999)
|
|
|
|
elif schedule == "sqrt_linear":
|
|
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
|
|
elif schedule == "sqrt":
|
|
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
|
|
else:
|
|
raise ValueError(f"schedule '{schedule}' unknown.")
|
|
return betas.numpy()
|
|
|
|
|
|
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
|
|
# select alphas for computing the variance schedule
|
|
alphas = alphacums[ddim_timesteps]
|
|
alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
|
|
|
# according the the formula provided in https://arxiv.org/abs/2010.02502
|
|
sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
|
|
if verbose:
|
|
print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
|
|
print(f'For the chosen value of eta, which is {eta}, '
|
|
f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
|
|
return sigmas, alphas, alphas_prev
|
|
|
|
|
|
def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
|
|
if ddim_discr_method == 'uniform':
|
|
c = num_ddpm_timesteps // num_ddim_timesteps
|
|
ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
|
|
elif ddim_discr_method == 'quad':
|
|
ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
|
|
else:
|
|
raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
|
|
|
|
# assert ddim_timesteps.shape[0] == num_ddim_timesteps
|
|
# add one to get the final alpha values right (the ones from first scale to data during sampling)
|
|
steps_out = ddim_timesteps + 1
|
|
if verbose:
|
|
print(f'Selected timesteps for ddim sampler: {steps_out}')
|
|
return steps_out
|
|
|
|
|
|
def noise_like(shape, device, repeat=False):
|
|
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
|
|
noise = lambda: torch.randn(shape, device=device)
|
|
return repeat_noise() if repeat else noise()
|
|
|
|
|
|
def timestep_embedding(device, timesteps, dim, max_period=10000, repeat_only=False):
|
|
"""
|
|
Create sinusoidal timestep embeddings.
|
|
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
|
These may be fractional.
|
|
:param dim: the dimension of the output.
|
|
:param max_period: controls the minimum frequency of the embeddings.
|
|
:return: an [N x dim] Tensor of positional embeddings.
|
|
"""
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
|
).to(device=device)
|
|
|
|
args = timesteps[:, None].float() * freqs[None]
|
|
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
|
return embedding
|