nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

replaced numpy ops with torch ops (#9)

Browse source 
* replaced numpy ops with torch ops

* polish

* polish
This commit is contained in:
Frank Lee 2024-03-30 13:01:58 +08:00 committed by GitHub
parent 223cc34c46
commit a0bdaced4e
5 changed files with 535 additions and 178 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
opensora/schedulers/iddpm/diffusion_utils.py
View file

@ -12,7 +12,7 @@
import numpy as np
import torch as th
import torch
def normal_kl(mean1, logvar1, mean2, logvar2):
@ -23,16 +23,18 @@ def normal_kl(mean1, logvar1, mean2, logvar2):
"""
tensor = None
for obj in (mean1, logvar1, mean2, logvar2):
if isinstance(obj, th.Tensor):
if isinstance(obj, torch.Tensor):
tensor = obj
break
assert tensor is not None, "at least one argument must be a Tensor"
# Force variances to be Tensors. Broadcasting helps convert scalars to
# Tensors, but it does not work for th.exp().
logvar1, logvar2 = [x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) for x in (logvar1, logvar2)]
# Tensors, but it does not work for torch.exp().
logvar1, logvar2 = [x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) for x in (logvar1, logvar2)]
return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * th.exp(-logvar2))
return 0.5 * (
-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
)
def approx_standard_normal_cdf(x):
@ -40,7 +42,7 @@ def approx_standard_normal_cdf(x):
A fast approximation of the cumulative distribution function of the
standard normal.
"""
return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
return 0.5 * (1.0 + torch.tanh(np.sqrt(2.0 / torch.pi) * (x + 0.044715 * torch.pow(x, 3))))
def continuous_gaussian_log_likelihood(x, *, means, log_scales):
@ -52,9 +54,9 @@ def continuous_gaussian_log_likelihood(x, *, means, log_scales):
:return: a tensor like x of log probabilities (in nats).
"""
centered_x = x - means
inv_stdv = th.exp(-log_scales)
inv_stdv = torch.exp(-log_scales)
normalized_x = centered_x * inv_stdv
log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x)
log_probs = torch.distributions.Normal(torch.zeros_like(x), torch.ones_like(x)).log_prob(normalized_x)
return log_probs
@ -70,18 +72,18 @@ def discretized_gaussian_log_likelihood(x, *, means, log_scales):
"""
assert x.shape == means.shape == log_scales.shape
centered_x = x - means
inv_stdv = th.exp(-log_scales)
inv_stdv = torch.exp(-log_scales)
plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
cdf_plus = approx_standard_normal_cdf(plus_in)
min_in = inv_stdv * (centered_x - 1.0 / 255.0)
cdf_min = approx_standard_normal_cdf(min_in)
log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
log_one_minus_cdf_min = torch.log((1.0 - cdf_min).clamp(min=1e-12))
cdf_delta = cdf_plus - cdf_min
log_probs = th.where(
log_probs = torch.where(
x < -0.999,
log_cdf_plus,
th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
torch.where(x > 0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))),
)
assert log_probs.shape == x.shape
return log_probs

182
opensora/schedulers/iddpm/gaussian_diffusion.py
View file

@ -11,16 +11,16 @@
# --------------------------------------------------------
import enum
import math
from typing import Callable, List
import numpy as np
import torch as th
import torch
from einops import rearrange
from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl
def mean_flat(tensor, mask=None):
def mean_flat(tensor: torch.Tensor, mask=None):
"""
Take the mean over all non-batch dimensions.
"""
@ -68,44 +68,65 @@ class LossType(enum.Enum):
return self == LossType.KL or self == LossType.RESCALED_KL
def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
def _warmup_beta(beta_start: float, beta_end: float, num_diffusion_timesteps: int, warmup_frac: float) -> torch.Tensor:
betas = beta_end * torch.ones(num_diffusion_timesteps, dtype=torch.float64)
warmup_time = int(num_diffusion_timesteps * warmup_frac)
betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
betas[:warmup_time] = torch.linspace(beta_start, beta_end, warmup_time, dtype=torch.float64)
return betas
def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
def get_beta_schedule(
beta_schedule: str, *, beta_start: float, beta_end: float, num_diffusion_timesteps: int
) -> torch.Tensor:
"""
This is the deprecated API for creating beta schedules.
See get_named_beta_schedule() for the new library of schedules.
"""
if beta_schedule == "quad":
betas = (
np.linspace(
torch.linspace(
beta_start**0.5,
beta_end**0.5,
num_diffusion_timesteps,
dtype=np.float64,
dtype=torch.float64,
)
** 2
)
elif beta_schedule == "linear":
betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
betas = torch.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=torch.float64)
elif beta_schedule == "warmup10":
betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1)
elif beta_schedule == "warmup50":
betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5)
elif beta_schedule == "const":
betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
betas = beta_end * torch.ones(num_diffusion_timesteps, dtype=torch.float64)
elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1
betas = 1.0 / np.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64)
betas = 1.0 / torch.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=torch.float64)
else:
raise NotImplementedError(beta_schedule)
assert betas.shape == (num_diffusion_timesteps,)
return betas
def betas_for_alpha_bar(num_diffusion_timesteps: int, alpha_bar: Callable, max_beta: float = 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].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.
"""
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 torch.DoubleTensor(betas)
def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
"""
Get a pre-defined beta schedule for the given name.
@ -127,31 +148,12 @@ def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
elif schedule_name == "squaredcos_cap_v2":
return betas_for_alpha_bar(
num_diffusion_timesteps,
lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
lambda t: matorch.cos((t + 0.008) / 1.008 * matorch.pi / 2) ** 2,
)
else:
raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, 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].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.
"""
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)
class GaussianDiffusion:
"""
Utilities for training and sampling diffusion models.
@ -161,43 +163,51 @@ class GaussianDiffusion:
starting at T and going to 1.
"""
def __init__(self, *, betas, model_mean_type, model_var_type, loss_type):
def __init__(
self, *, betas: torch.Tensor, model_mean_type: str, model_var_type: str, loss_type: str, device: str = "cuda"
):
if device == "cuda":
device = torch.device(f"cuda:{torch.cuda.current_device()}")
elif device == "cpu":
device = torch.device("cpu")
else:
raise ValueError(f"Unknown device: {device}")
self.device = device
self.model_mean_type = model_mean_type
self.model_var_type = model_var_type
self.loss_type = loss_type
# Use float64 for accuracy.
betas = np.array(betas, dtype=np.float64)
self.betas = betas
assert len(betas.shape) == 1, "betas must be 1-D"
assert (betas > 0).all() and (betas <= 1).all()
self.betas = betas.to(self.device)
assert len(self.betas.shape) == 1, "betas must be 1-D"
assert (self.betas > 0).all() and (self.betas <= 1).all()
self.num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
self.alphas_cumprod = np.cumprod(alphas, axis=0)
self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(alphas, axis=0)
self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0], device=self.device), self.alphas_cumprod[:-1]])
self.alphas_cumprod_next = torch.cat([self.alphas_cumprod[1:], torch.tensor([0.0], device=self.device)])
assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
# calculations for diffusion q(x_t | x_{t-1}) and others
self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
self.log_one_minus_alphas_cumprod = torch.log(1.0 - self.alphas_cumprod)
self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)
self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)
# calculations for posterior q(x_{t-1} | x_t, x_0)
self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
self.posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
self.posterior_log_variance_clipped = (
np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))
torch.log(torch.cat([self.posterior_variance[1].unsqueeze(0), self.posterior_variance[1:]]))
if len(self.posterior_variance) > 1
else np.array([])
else torch.DoubleTensor([])
)
self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
self.posterior_mean_coef1 = self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - self.alphas_cumprod)
def q_mean_variance(self, x_start, t):
"""
@ -221,7 +231,7 @@ class GaussianDiffusion:
:return: A noisy version of x_start.
"""
if noise is None:
noise = th.randn_like(x_start)
noise = torch.randn_like(x_start)
assert noise.shape == x_start.shape
return (
_extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
@ -281,20 +291,20 @@ class GaussianDiffusion:
if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
assert model_output.shape == (B, C * 2, *x.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
model_output, model_var_values = torch.split(model_output, C, dim=1)
min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)
max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
max_log = _extract_into_tensor(torch.log(self.betas), t, x.shape)
# The model_var_values is [-1, 1] for [min_var, max_var].
frac = (model_var_values + 1) / 2
model_log_variance = frac * max_log + (1 - frac) * min_log
model_variance = th.exp(model_log_variance)
model_variance = torch.exp(model_log_variance)
else:
model_variance, model_log_variance = {
# for fixedlarge, we set the initial (log-)variance like so
# to get a better decoder log likelihood.
ModelVarType.FIXED_LARGE: (
np.append(self.posterior_variance[1], self.betas[1:]),
np.log(np.append(self.posterior_variance[1], self.betas[1:])),
torch.cat(self.posterior_variance[1].unsqueeze(0), self.betas[1:]),
torch.log(torch.cat(self.posterior_variance[1].unsqueeze(0), self.betas[1:])),
),
ModelVarType.FIXED_SMALL: (
self.posterior_variance,
@ -403,15 +413,15 @@ class GaussianDiffusion:
denoised_fn=denoised_fn,
model_kwargs=model_kwargs,
)
noise = th.randn_like(x)
noise = torch.randn_like(x)
nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) # no noise when t == 0
if cond_fn is not None:
out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
sample = out["mean"] + nonzero_mask * torch.exp(0.5 * out["log_variance"]) * noise
if mask is not None:
if mask.shape[0] != x.shape[0]:
mask = mask.repeat(2, 1) # HACK
sample = th.where(mask[:, None, :, None, None], sample, x)
sample = torch.where(mask[:, None, :, None, None], sample, x)
return {"sample": sample, "pred_xstart": out["pred_xstart"]}
@ -488,7 +498,7 @@ class GaussianDiffusion:
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device)
img = torch.randn(*shape, device=device)
indices = list(range(self.num_timesteps))[::-1]
if progress:
@ -498,8 +508,8 @@ class GaussianDiffusion:
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
t = torch.tensor([i] * shape[0], device=device)
with torch.no_grad():
out = self.p_sample(
model,
img,
@ -545,10 +555,10 @@ class GaussianDiffusion:
alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(1 - alpha_bar / alpha_bar_prev)
sigma = eta * torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * torch.sqrt(1 - alpha_bar / alpha_bar_prev)
# Equation 12.
noise = th.randn_like(x)
mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
noise = torch.randn_like(x)
mean_pred = out["pred_xstart"] * torch.sqrt(alpha_bar_prev) + torch.sqrt(1 - alpha_bar_prev - sigma**2) * eps
nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) # no noise when t == 0
sample = mean_pred + nonzero_mask * sigma * noise
return {"sample": sample, "pred_xstart": out["pred_xstart"]}
@ -586,7 +596,7 @@ class GaussianDiffusion:
alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
# Equation 12. reversed
mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps
mean_pred = out["pred_xstart"] * torch.sqrt(alpha_bar_next) + torch.sqrt(1 - alpha_bar_next) * eps
return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
@ -647,7 +657,7 @@ class GaussianDiffusion:
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device)
img = torch.randn(*shape, device=device)
indices = list(range(self.num_timesteps))[::-1]
if progress:
@ -657,8 +667,8 @@ class GaussianDiffusion:
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
t = torch.tensor([i] * shape[0], device=device)
with torch.no_grad():
out = self.ddim_sample(
model,
img,
@ -694,7 +704,7 @@ class GaussianDiffusion:
# At the first timestep return the decoder NLL,
# otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
output = th.where((t == 0), decoder_nll, kl)
output = torch.where((t == 0), decoder_nll, kl)
return {"output": output, "pred_xstart": out["pred_xstart"]}
def training_losses(self, model, x_start, t, model_kwargs=None, noise=None, mask=None, weights=None):
@ -712,12 +722,12 @@ class GaussianDiffusion:
if model_kwargs is None:
model_kwargs = {}
if noise is None:
noise = th.randn_like(x_start)
noise = torch.randn_like(x_start)
x_t = self.q_sample(x_start, t, noise=noise)
if mask is not None:
t0 = th.zeros_like(t)
t0 = torch.zeros_like(t)
x_t0 = self.q_sample(x_start, t0, noise=noise)
x_t = th.where(mask[:, None, :, None, None], x_t, x_t0)
x_t = torch.where(mask[:, None, :, None, None], x_t, x_t0)
terms = {}
@ -742,10 +752,10 @@ class GaussianDiffusion:
]:
B, C = x_t.shape[:2]
assert model_output.shape == (B, C * 2, *x_t.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
model_output, model_var_values = torch.split(model_output, C, dim=1)
# Learn the variance using the variational bound, but don't let
# it affect our mean prediction.
frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
frozen_out = torch.cat([model_output.detach(), model_var_values], dim=1)
terms["vb"] = self._vb_terms_bpd(
model=lambda *args, r=frozen_out: r,
x_start=x_start,
@ -788,7 +798,7 @@ class GaussianDiffusion:
:return: a batch of [N] KL values (in bits), one per batch element.
"""
batch_size = x_start.shape[0]
t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
return mean_flat(kl_prior) / np.log(2.0)
@ -816,11 +826,11 @@ class GaussianDiffusion:
xstart_mse = []
mse = []
for t in list(range(self.num_timesteps))[::-1]:
t_batch = th.tensor([t] * batch_size, device=device)
noise = th.randn_like(x_start)
t_batch = torch.tensor([t] * batch_size, device=device)
noise = torch.randn_like(x_start)
x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
# Calculate VLB term at the current timestep
with th.no_grad():
with torch.no_grad():
out = self._vb_terms_bpd(
model,
x_start=x_start,
@ -834,9 +844,9 @@ class GaussianDiffusion:
eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
mse.append(mean_flat((eps - noise) ** 2))
vb = th.stack(vb, dim=1)
xstart_mse = th.stack(xstart_mse, dim=1)
mse = th.stack(mse, dim=1)
vb = torch.stack(vb, dim=1)
xstart_mse = torch.stack(xstart_mse, dim=1)
mse = torch.stack(mse, dim=1)
prior_bpd = self._prior_bpd(x_start)
total_bpd = vb.sum(dim=1) + prior_bpd
@ -849,7 +859,7 @@ class GaussianDiffusion:
}
def _extract_into_tensor(arr, timesteps, broadcast_shape):
def _extract_into_tensor(arr: torch.Tensor, timesteps: torch.Tensor, broadcast_shape: List[int]):
"""
Extract values from a 1-D numpy array for a batch of indices.
:param arr: the 1-D numpy array.
@ -858,7 +868,7 @@ def _extract_into_tensor(arr, timesteps, broadcast_shape):
dimension equal to the length of timesteps.
:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
"""
res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
res = arr.to(timesteps.device)[timesteps].float()
while len(res.shape) < len(broadcast_shape):
res = res[..., None]
return res + th.zeros(broadcast_shape, device=timesteps.device)
return res + torch.zeros(broadcast_shape, device=timesteps.device)

7
opensora/schedulers/iddpm/respace.py
View file

@ -11,8 +11,7 @@
# --------------------------------------------------------
import numpy as np
import torch as th
import torch
from .gaussian_diffusion import GaussianDiffusion
@ -87,7 +86,7 @@ class SpacedDiffusion(GaussianDiffusion):
new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
last_alpha_cumprod = alpha_cumprod
self.timestep_map.append(i)
kwargs["betas"] = np.array(new_betas)
kwargs["betas"] = torch.FloatTensor(new_betas)
super().__init__(**kwargs)
def p_mean_variance(self, model, *args, **kwargs): # pylint: disable=signature-differs
@ -120,7 +119,7 @@ class _WrappedModel:
self.original_num_steps = original_num_steps
def __call__(self, x, ts, **kwargs):
map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
map_tensor = torch.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
new_ts = map_tensor[ts]
# if self.rescale_timesteps:
# new_ts = new_ts.float() * (1000.0 / self.original_num_steps)

150
opensora/schedulers/iddpm/speed.py
View file

@ -1,75 +1,75 @@
import numpy as np
import torch
import torch.nn.functional as F
from opensora.registry import SCHEDULERS
from . import gaussian_diffusion as gd
from .respace import SpacedDiffusion, space_timesteps
@SCHEDULERS.register_module("iddpm-speed")
class SpeeDiffusion(SpacedDiffusion):
def __init__(
self,
num_sampling_steps=None,
timestep_respacing=None,
noise_schedule="linear",
use_kl=False,
sigma_small=False,
predict_xstart=False,
learn_sigma=True,
rescale_learned_sigmas=False,
diffusion_steps=1000,
cfg_scale=4.0,
):
betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
if use_kl:
loss_type = gd.LossType.RESCALED_KL
elif rescale_learned_sigmas:
loss_type = gd.LossType.RESCALED_MSE
else:
loss_type = gd.LossType.MSE
if num_sampling_steps is not None:
assert timestep_respacing is None
timestep_respacing = str(num_sampling_steps)
if timestep_respacing is None or timestep_respacing == "":
timestep_respacing = [diffusion_steps]
super().__init__(
use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
betas=betas,
model_mean_type=(gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X),
model_var_type=(
(gd.ModelVarType.FIXED_LARGE if not sigma_small else gd.ModelVarType.FIXED_SMALL)
if not learn_sigma
else gd.ModelVarType.LEARNED_RANGE
),
loss_type=loss_type,
)
self.cfg_scale = cfg_scale
grad = np.gradient(self.sqrt_one_minus_alphas_cumprod)
self.meaningful_steps = np.argmax(grad < 5e-5) + 1
# p2 weighting from: Perception Prioritized Training of Diffusion Models
self.p2_gamma = 1
self.p2_k = 1
self.snr = 1.0 / (1 - self.alphas_cumprod) - 1
sqrt_one_minus_alphas_bar = torch.from_numpy(self.sqrt_one_minus_alphas_cumprod)
p = torch.tanh(1e6 * (torch.gradient(sqrt_one_minus_alphas_bar)[0] - 1e-4)) + 1.5
self.p = F.normalize(p, p=1, dim=0)
self.weights = 1 / (self.p2_k + self.snr) ** self.p2_gamma
def t_sample(self, n, device):
t = torch.multinomial(self.p, n // 2 + 1, replacement=True).to(device)
dual_t = torch.where(t < self.meaningful_steps, self.meaningful_steps - t, t - self.meaningful_steps)
t = torch.cat([t, dual_t], dim=0)[:n]
return t
def training_losses(self, model, x, t, *args, **kwargs): # pylint: disable=signature-differs
t = self.t_sample(x.shape[0], x.device)
return super().training_losses(model, x, t, weights=self.weights, *args, **kwargs)
def sample(self, *args, **kwargs):
raise NotImplementedError("SpeeDiffusion is only for training")
import numpy as np
import torch
import torch.nn.functional as F
from opensora.registry import SCHEDULERS
from . import gaussian_diffusion as gd
from .respace import SpacedDiffusion, space_timesteps
@SCHEDULERS.register_module("iddpm-speed")
class SpeeDiffusion(SpacedDiffusion):
def __init__(
self,
num_sampling_steps=None,
timestep_respacing=None,
noise_schedule="linear",
use_kl=False,
sigma_small=False,
predict_xstart=False,
learn_sigma=True,
rescale_learned_sigmas=False,
diffusion_steps=1000,
cfg_scale=4.0,
):
betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
if use_kl:
loss_type = gd.LossType.RESCALED_KL
elif rescale_learned_sigmas:
loss_type = gd.LossType.RESCALED_MSE
else:
loss_type = gd.LossType.MSE
if num_sampling_steps is not None:
assert timestep_respacing is None
timestep_respacing = str(num_sampling_steps)
if timestep_respacing is None or timestep_respacing == "":
timestep_respacing = [diffusion_steps]
super().__init__(
use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
betas=betas,
model_mean_type=(gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X),
model_var_type=(
(gd.ModelVarType.FIXED_LARGE if not sigma_small else gd.ModelVarType.FIXED_SMALL)
if not learn_sigma
else gd.ModelVarType.LEARNED_RANGE
),
loss_type=loss_type,
)
self.cfg_scale = cfg_scale
# we fallback to numpy here as argmax_cuda is not implemented for Bool
grad = np.gradient(self.sqrt_one_minus_alphas_cumprod.cpu())
self.meaningful_steps = np.argmax(grad < 5e-5) + 1
# p2 weighting from: Perception Prioritized Training of Diffusion Models
self.p2_gamma = 1
self.p2_k = 1
self.snr = 1.0 / (1 - self.alphas_cumprod) - 1
sqrt_one_minus_alphas_bar = self.sqrt_one_minus_alphas_cumprod
p = torch.tanh(1e6 * (torch.gradient(sqrt_one_minus_alphas_bar)[0] - 1e-4)) + 1.5
self.p = F.normalize(p, p=1, dim=0)
self.weights = 1 / (self.p2_k + self.snr) ** self.p2_gamma
def t_sample(self, n, device):
t = torch.multinomial(self.p, n // 2 + 1, replacement=True).to(device)
dual_t = torch.where(t < self.meaningful_steps, self.meaningful_steps - t, t - self.meaningful_steps)
t = torch.cat([t, dual_t], dim=0)[:n]
return t
def training_losses(self, model, x, t, *args, **kwargs): # pylint: disable=signature-differs
t = self.t_sample(x.shape[0], x.device)
return super().training_losses(model, x, t, weights=self.weights, *args, **kwargs)
def sample(self, *args, **kwargs):
raise NotImplementedError("SpeeDiffusion is only for training")

346
tests/test_np_torch.py Normal file
View file

@ -0,0 +1,346 @@
from typing import Callable
import numpy as np
import torch
# ==================================
# Warm Up Beta
# ==================================
def _warmup_beta_numpy(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
warmup_time = int(num_diffusion_timesteps * warmup_frac)
betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
return betas
def _warmup_beta_torch(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
betas = beta_end * torch.ones(num_diffusion_timesteps, dtype=torch.float64)
warmup_time = int(num_diffusion_timesteps * warmup_frac)
betas[:warmup_time] = torch.linspace(beta_start, beta_end, warmup_time, dtype=torch.float64)
return betas
def test_warmup_beta():
beta_start = 1e-6
beta_end = 0.99
num_diffusion_timesteps = 1000
warmup_frac = 0.1
betas_np = _warmup_beta_numpy(beta_start, beta_end, num_diffusion_timesteps, warmup_frac)
betas_torch = _warmup_beta_torch(beta_start, beta_end, num_diffusion_timesteps, warmup_frac)
assert np.allclose(betas_np, betas_torch.numpy())
print("Test passed for warmup_beta()")
# ==================================
# Beta Schedule
# ==================================
def get_beta_schedule_numpy(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
"""
This is the deprecated API for creating beta schedules.
See get_named_beta_schedule() for the new library of schedules.
"""
if beta_schedule == "quad":
betas = (
np.linspace(
beta_start**0.5,
beta_end**0.5,
num_diffusion_timesteps,
dtype=np.float64,
)
** 2
)
elif beta_schedule == "linear":
betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
elif beta_schedule == "warmup10":
betas = _warmup_beta_numpy(beta_start, beta_end, num_diffusion_timesteps, 0.1)
elif beta_schedule == "warmup50":
betas = _warmup_beta_numpy(beta_start, beta_end, num_diffusion_timesteps, 0.5)
elif beta_schedule == "const":
betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1
betas = 1.0 / np.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64)
else:
raise NotImplementedError(beta_schedule)
assert betas.shape == (num_diffusion_timesteps,)
return betas
def get_beta_schedule_torch(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
"""
This is the deprecated API for creating beta schedules.
See get_named_beta_schedule() for the new library of schedules.
"""
if beta_schedule == "quad":
betas = (
np.linspace(
beta_start**0.5,
beta_end**0.5,
num_diffusion_timesteps,
dtype=np.float64,
)
** 2
)
elif beta_schedule == "linear":
betas = torch.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=torch.float64)
elif beta_schedule == "warmup10":
betas = _warmup_beta_torch(beta_start, beta_end, num_diffusion_timesteps, 0.1)
elif beta_schedule == "warmup50":
betas = _warmup_beta_torch(beta_start, beta_end, num_diffusion_timesteps, 0.5)
elif beta_schedule == "const":
betas = beta_end * torch.ones(num_diffusion_timesteps, dtype=np.float64)
elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1
betas = 1.0 / torch.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=torch.float64)
else:
raise NotImplementedError(beta_schedule)
assert betas.shape == (num_diffusion_timesteps,)
return betas
def test_get_beta_Schedule():
beta_start = 1e-6
beta_end = 0.99
num_diffusion_timesteps = 1000
beta_schedule = "linear"
betas_np = get_beta_schedule_numpy(
beta_schedule, beta_start=beta_start, beta_end=beta_end, num_diffusion_timesteps=num_diffusion_timesteps
)
betas_torch = get_beta_schedule_torch(
beta_schedule, beta_start=beta_start, beta_end=beta_end, num_diffusion_timesteps=num_diffusion_timesteps
)
assert np.allclose(betas_np, betas_torch.numpy())
print("Test passed for get_beta_schedule()")
# ====================
# Replace alpha
# ====================
def betas_for_alpha_bar_numpy(num_diffusion_timesteps: int, alpha_bar: Callable, max_beta: float = 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].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.
"""
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)
def betas_for_alpha_bar_torch(num_diffusion_timesteps: int, alpha_bar: Callable, max_beta: float = 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].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.
"""
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 torch.DoubleTensor(betas)
def test_betas_for_alpha_bar():
num_diffusion_timesteps = 1000
alpha_bar = lambda t: 1 - t
max_beta = 0.999
betas_np = betas_for_alpha_bar_numpy(num_diffusion_timesteps, alpha_bar, max_beta)
betas_torch = betas_for_alpha_bar_torch(num_diffusion_timesteps, alpha_bar, max_beta)
assert np.allclose(betas_np, betas_torch.numpy())
print("Test passed for betas_for_alpha_bar()")
# =======================
# Gaussian init
# =======================
def init_numpy(betas):
# Use float64 for accuracy.
betas = torch.DoubleTensor(betas)
assert len(betas.shape) == 1, "betas must be 1-D"
assert (betas > 0).all() and (betas <= 1).all()
num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])
alphas_cumprod_next = np.append(alphas_cumprod[1:], 0.0)
assert alphas_cumprod_prev.shape == (num_timesteps,)
# calculations for diffusion q(x_t | x_{t-1}) and others
np.sqrt(alphas_cumprod)
np.sqrt(1.0 - alphas_cumprod)
np.log(1.0 - alphas_cumprod)
np.sqrt(1.0 / alphas_cumprod)
np.sqrt(1.0 / alphas_cumprod - 1)
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
posterior_log_variance_clipped = (
np.log(np.append(posterior_variance[1], posterior_variance[1:]))
if len(posterior_variance) > 1
else np.array([])
)
posterior_mean_coef1 = betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
posterior_mean_coef2 = (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod)
return alphas_cumprod_prev, alphas_cumprod_next, posterior_mean_coef1, posterior_mean_coef2
def gaussian_init_numpy(betas):
# Use float64 for accuracy.
betas = np.array(betas, dtype=np.float64)
assert len(betas.shape) == 1, "betas must be 1-D"
assert (betas > 0).all() and (betas <= 1).all()
num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])
alphas_cumprod_next = np.append(alphas_cumprod[1:], 0.0)
assert alphas_cumprod_prev.shape == (num_timesteps,)
# calculations for diffusion q(x_t | x_{t-1}) and others
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
log_one_minus_alphas_cumprod = np.log(1.0 - alphas_cumprod)
sqrt_recip_alphas_cumprod = np.sqrt(1.0 / alphas_cumprod)
sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / alphas_cumprod - 1)
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
posterior_log_variance_clipped = (
np.log(np.append(posterior_variance[1], posterior_variance[1:]))
if len(posterior_variance) > 1
else np.array([])
)
posterior_mean_coef1 = betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
posterior_mean_coef2 = (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod)
return (
alphas_cumprod_prev,
alphas_cumprod_next,
sqrt_alphas_cumprod,
sqrt_one_minus_alphas_cumprod,
log_one_minus_alphas_cumprod,
sqrt_recip_alphas_cumprod,
sqrt_recipm1_alphas_cumprod,
posterior_log_variance_clipped,
posterior_mean_coef1,
posterior_mean_coef2,
)
def gaussian_init_torch(betas):
# Use float64 for accuracy.
betas = torch.DoubleTensor(betas)
assert len(betas.shape) == 1, "betas must be 1-D"
assert (betas > 0).all() and (betas <= 1).all()
num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, axis=0)
alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])
alphas_cumprod_next = torch.cat([alphas_cumprod[1:], torch.tensor([0.0])])
assert alphas_cumprod_prev.shape == (num_timesteps,)
# calculations for diffusion q(x_t | x_{t-1}) and others
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
log_one_minus_alphas_cumprod = torch.log(1.0 - alphas_cumprod)
sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod)
sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod - 1)
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
posterior_log_variance_clipped = (
torch.log(torch.cat([posterior_variance[1].unsqueeze(0), posterior_variance[1:]]))
if len(posterior_variance) > 1
else torch.array([])
)
posterior_mean_coef1 = betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
posterior_mean_coef2 = (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod)
return (
alphas_cumprod_prev,
alphas_cumprod_next,
sqrt_alphas_cumprod,
sqrt_one_minus_alphas_cumprod,
log_one_minus_alphas_cumprod,
sqrt_recip_alphas_cumprod,
sqrt_recipm1_alphas_cumprod,
posterior_log_variance_clipped,
posterior_mean_coef1,
posterior_mean_coef2,
)
def test_gaussian_init():
betas = np.linspace(1e-6, 0.99, 1000)
(
alphas_cumprod_prev,
alphas_cumprod_next,
sqrt_alphas_cumprod,
sqrt_one_minus_alphas_cumprod,
log_one_minus_alphas_cumprod,
sqrt_recip_alphas_cumprod,
sqrt_recipm1_alphas_cumprod,
posterior_log_variance_clipped,
posterior_mean_coef1,
posterior_mean_coef2,
) = gaussian_init_numpy(betas)
(
alphas_cumprod_prev_t,
alphas_cumprod_next_t,
sqrt_alphas_cumprod_t,
sqrt_one_minus_alphas_cumprod_t,
log_one_minus_alphas_cumprod_t,
sqrt_recip_alphas_cumprod_t,
sqrt_recipm1_alphas_cumprod_t,
posterior_log_variance_clipped_t,
posterior_mean_coef1_t,
posterior_mean_coef2_t,
) = gaussian_init_torch(betas)
assert np.allclose(alphas_cumprod_prev, alphas_cumprod_prev_t.numpy())
assert np.allclose(alphas_cumprod_next, alphas_cumprod_next_t.numpy())
assert np.allclose(sqrt_alphas_cumprod, sqrt_alphas_cumprod_t.numpy())
assert np.allclose(sqrt_one_minus_alphas_cumprod, sqrt_one_minus_alphas_cumprod_t.numpy())
assert np.allclose(log_one_minus_alphas_cumprod, log_one_minus_alphas_cumprod_t.numpy())
assert np.allclose(sqrt_recip_alphas_cumprod, sqrt_recip_alphas_cumprod_t.numpy())
assert np.allclose(sqrt_recipm1_alphas_cumprod, sqrt_recipm1_alphas_cumprod_t.numpy())
assert np.allclose(posterior_log_variance_clipped, posterior_log_variance_clipped_t.numpy())
assert np.allclose(posterior_mean_coef1, posterior_mean_coef1_t.numpy())
assert np.allclose(posterior_mean_coef2, posterior_mean_coef2_t.numpy())
print("Test passed for gaussian_init()")
if __name__ == "__main__":
test_warmup_beta()
test_get_beta_Schedule()
test_betas_for_alpha_bar()
test_gaussian_init()