File size: 5,647 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | from typing import Literal
import math
import torch
def make_beta_schedule(
schedule: Literal["cosine", "sigmoid", "sd", "linear", "alphas_cumprod_linear"],
shift: float = 1.0,
clip_min: float = 1e-9,
zero_terminal_snr: bool = True,
**kwargs,
):
schedule_fn = {
"alphas_cumprod_linear": alphas_cumprod_linear_schedule,
"cosine": cosine_schedule,
"cosine_simple_diffusion": cosine_simple_diffusion_schedule,
"sigmoid": sigmoid_schedule,
"sd": sd_schedule,
"linear": beta_linear_schedule,
}[schedule]
alphas_cumprod = schedule_fn(**kwargs)
if schedule not in ["cosine", "cosine_simple_diffusion"] and zero_terminal_snr:
# cosine schedule already enforces zero terminal SNR
# simple diffusion's cosine schedule shall not enforce zero terminal SNR
alphas_cumprod = enforce_zero_terminal_snr(alphas_cumprod)
if (
shift != 1.0 and schedule != "cosine_simple_diffusion"
): # cosine_simple_diffusion already has shift built in
alphas_cumprod = shift_beta_schedule(alphas_cumprod, shift)
alphas = alphas_cumprod[1:] / alphas_cumprod[:-1]
alphas = torch.cat([alphas_cumprod[0:1], alphas])
betas = 1 - alphas
return torch.clip(betas, clip_min, 1.0)
def cosine_schedule(timesteps, s=0.008):
"""
cosine schedule
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
"""
steps = timesteps + 1
t = torch.linspace(0, timesteps, steps, dtype=torch.float64) / timesteps
alphas_cumprod = torch.cos((t + s) / (1 + s) * math.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
return alphas_cumprod[1:]
def cosine_simple_diffusion_schedule(
timesteps,
logsnr_min=-15.0,
logsnr_max=15.0,
shifted: float = 1.0,
interpolated: bool = False,
):
"""
cosine schedule with different parameterization
following Simple Diffusion - https://arxiv.org/abs/2301.11093
Supports "shifted cosine schedule" and "interpolated cosine schedule"
Args:
timesteps: number of timesteps
logsnr_min: minimum log SNR
logsnr_max: maximum log SNR
shifted: shift the schedule by a factor. Should be base_resolution / current_resolution
interpolated: interpolate between the original and the shifted schedule, requires shifted != 1.0
"""
t_min = torch.atan(torch.exp(-0.5 * torch.tensor(logsnr_max, dtype=torch.float64)))
t_max = torch.atan(torch.exp(-0.5 * torch.tensor(logsnr_min, dtype=torch.float64)))
t = torch.linspace(0, 1, timesteps, dtype=torch.float64)
logsnr = -2 * torch.log(torch.tan(t_min + t * (t_max - t_min)))
if shifted != 1.0:
shifted_logsnr = logsnr + 2 * torch.log(
torch.tensor(shifted, dtype=torch.float64)
)
if interpolated:
logsnr = t * logsnr + (1 - t) * shifted_logsnr
else:
logsnr = shifted_logsnr
alphas_cumprod = 1 / (1 + torch.exp(-logsnr))
return alphas_cumprod
def alphas_cumprod_linear_schedule(timesteps: int) -> torch.Tensor:
"""
linear schedule
as proposed in https://arxiv.org/abs/2301.10972
"""
t = torch.linspace(0, timesteps, timesteps + 1, dtype=torch.float64) / timesteps
return (1 - t)[1:]
def beta_linear_schedule(
timesteps: int, start: float = 0.0001, end: float = 0.02
) -> torch.Tensor:
"""
linear schedule
as proposed in https://arxiv.org/abs/2006.11239 (original DDPM paper)
"""
betas = torch.linspace(start, end, timesteps, dtype=torch.float64)
return (1 - betas).cumprod(dim=0)
def sigmoid_schedule(timesteps, start=-3, end=3, tau=1):
"""
sigmoid schedule
proposed in https://arxiv.org/abs/2212.11972 - Figure 8
better for images > 64x64, when used during training
"""
steps = timesteps + 1
t = torch.linspace(0, timesteps, steps, dtype=torch.float64) / timesteps
v_start = torch.tensor(start / tau).sigmoid()
v_end = torch.tensor(end / tau).sigmoid()
alphas_cumprod = (-((t * (end - start) + start) / tau).sigmoid() + v_end) / (
v_end - v_start
)
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
return alphas_cumprod[1:]
def sd_schedule(timesteps, start=0.00085, end=0.0120):
"""
stable diffusion's noise schedule
https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/modules/diffusionmodules/util.py#L21
"""
betas = torch.linspace(start**0.5, end**0.5, timesteps, dtype=torch.float64) ** 2
alphas_cumprod = (1 - betas).cumprod(dim=0)
return alphas_cumprod
def shift_beta_schedule(alphas_cumprod: torch.Tensor, shift: float):
"""
scale alphas_cumprod so that SNR is multiplied by shift ** 2
"""
snr_scale = shift**2
return (snr_scale * alphas_cumprod) / (
snr_scale * alphas_cumprod + 1 - alphas_cumprod
)
def enforce_zero_terminal_snr(alphas_cumprod):
"""
enforce zero terminal SNR following https://arxiv.org/abs/2305.08891
returns betas
"""
alphas_cumprod_sqrt = torch.sqrt(alphas_cumprod)
# store old values
alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone()
alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone()
# shift so last timestep is zero
alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T
# scale so first timestep is back to original value
alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / alphas_cumprod_sqrt[0]
# convert to betas
alphas_cumprod = alphas_cumprod_sqrt**2
assert alphas_cumprod[-1] == 0, "terminal SNR not zero"
return alphas_cumprod
|