# everything that can improve v-prediction model
# dynamic scaling + tsnr + beta modifier + dynamic cfg rescale + ...
# written by lvmin at stanford 2024
import torch
import numpy as np
from tqdm import tqdm
from functools import partial
from diffusers_vdm.basics import extract_into_tensor
to_torch = partial(torch.tensor, dtype=torch.float32)
def rescale_zero_terminal_snr(betas):
# Convert betas to alphas_bar_sqrt
alphas = 1.0 - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_bar_sqrt = np.sqrt(alphas_cumprod)
# Store old values.
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].copy()
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].copy()
# Shift so the last timestep is zero.
alphas_bar_sqrt -= alphas_bar_sqrt_T
# Scale so the first timestep is back to the old value.
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
# Convert alphas_bar_sqrt to betas
alphas_bar = alphas_bar_sqrt**2 # Revert sqrt
alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod
alphas = np.concatenate([alphas_bar[0:1], alphas])
betas = 1 - alphas
return betas
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# rescale the results from guidance (fixes overexposure)
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
# mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
return noise_cfg
class SamplerDynamicTSNR(torch.nn.Module):
@torch.no_grad()
def __init__(self, unet, terminal_scale=0.7):
super().__init__()
self.unet = unet
self.is_v = True
self.n_timestep = 1000
self.guidance_rescale = 0.7
linear_start = 0.00085
linear_end = 0.012
betas = np.linspace(linear_start ** 0.5, linear_end ** 0.5, self.n_timestep, dtype=np.float64) ** 2
betas = rescale_zero_terminal_snr(betas)
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod).to(unet.device))
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)).to(unet.device))
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)).to(unet.device))
# Dynamic TSNR
turning_step = 400
scale_arr = np.concatenate([
np.linspace(1.0, terminal_scale, turning_step),
np.full(self.n_timestep - turning_step, terminal_scale)
])
self.register_buffer('scale_arr', to_torch(scale_arr).to(unet.device))
def predict_eps_from_z_and_v(self, x_t, t, v):
return self.sqrt_alphas_cumprod[t] * v + self.sqrt_one_minus_alphas_cumprod[t] * x_t
def predict_start_from_z_and_v(self, x_t, t, v):
return self.sqrt_alphas_cumprod[t] * x_t - self.sqrt_one_minus_alphas_cumprod[t] * v
def q_sample(self, x0, t, noise):
return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x0.shape) * x0 +
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
def get_v(self, x0, t, noise):
return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x0.shape) * noise -
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x0.shape) * x0)
def dynamic_x0_rescale(self, x0, t):
return x0 * extract_into_tensor(self.scale_arr, t, x0.shape)
@torch.no_grad()
def get_ground_truth(self, x0, noise, t):
x0 = self.dynamic_x0_rescale(x0, t)
xt = self.q_sample(x0, t, noise)
target = self.get_v(x0, t, noise) if self.is_v else noise
return xt, target
def get_uniform_trailing_steps(self, steps):
c = self.n_timestep / steps
ddim_timesteps = np.flip(np.round(np.arange(self.n_timestep, 0, -c))).astype(np.int64)
steps_out = ddim_timesteps - 1
return torch.tensor(steps_out, device=self.unet.device, dtype=torch.long)
@torch.no_grad()
def forward(self, latent_shape, steps, extra_args, progress_tqdm=None):
bar = tqdm if progress_tqdm is None else progress_tqdm
eta = 1.0
timesteps = self.get_uniform_trailing_steps(steps)
timesteps_prev = torch.nn.functional.pad(timesteps[:-1], pad=(1, 0))
x = torch.randn(latent_shape, device=self.unet.device, dtype=self.unet.dtype)
alphas = self.alphas_cumprod[timesteps]
alphas_prev = self.alphas_cumprod[timesteps_prev]
scale_arr = self.scale_arr[timesteps]
scale_arr_prev = self.scale_arr[timesteps_prev]
sqrt_one_minus_alphas = torch.sqrt(1 - alphas)
sigmas = eta * np.sqrt((1 - alphas_prev.cpu().numpy()) / (1 - alphas.cpu()) * (1 - alphas.cpu() / alphas_prev.cpu().numpy()))
s_in = x.new_ones((x.shape[0]))
s_x = x.new_ones((x.shape[0], ) + (1, ) * (x.ndim - 1))
for i in bar(range(len(timesteps))):
index = len(timesteps) - 1 - i
t = timesteps[index].item()
model_output = self.model_apply(x, t * s_in, **extra_args)
if self.is_v:
e_t = self.predict_eps_from_z_and_v(x, t, model_output)
else:
e_t = model_output
a_prev = alphas_prev[index].item() * s_x
sigma_t = sigmas[index].item() * s_x
if self.is_v:
pred_x0 = self.predict_start_from_z_and_v(x, t, model_output)
else:
a_t = alphas[index].item() * s_x
sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
# dynamic rescale
scale_t = scale_arr[index].item() * s_x
prev_scale_t = scale_arr_prev[index].item() * s_x
rescale = (prev_scale_t / scale_t)
pred_x0 = pred_x0 * rescale
dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
noise = sigma_t * torch.randn_like(x)
x = a_prev.sqrt() * pred_x0 + dir_xt + noise
return x
@torch.no_grad()
def model_apply(self, x, t, **extra_args):
x = x.to(device=self.unet.device, dtype=self.unet.dtype)
cfg_scale = extra_args['cfg_scale']
p = self.unet(x, t, **extra_args['positive'])
n = self.unet(x, t, **extra_args['negative'])
o = n + cfg_scale * (p - n)
o_better = rescale_noise_cfg(o, p, guidance_rescale=self.guidance_rescale)
return o_better