978 lines
29 KiB
Python
978 lines
29 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""
|
|
Based on: https://github.com/crowsonkb/k-diffusion
|
|
"""
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch as th
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# from piq import LPIPS
|
|
from utils.ssim import SSIM
|
|
|
|
from modules.diffusion.karras.random_utils import get_generator
|
|
|
|
|
|
def mean_flat(tensor):
|
|
"""
|
|
Take the mean over all non-batch dimensions.
|
|
"""
|
|
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
|
|
|
|
|
def append_dims(x, target_dims):
|
|
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
|
dims_to_append = target_dims - x.ndim
|
|
if dims_to_append < 0:
|
|
raise ValueError(
|
|
f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
|
|
)
|
|
return x[(...,) + (None,) * dims_to_append]
|
|
|
|
|
|
def append_zero(x):
|
|
return th.cat([x, x.new_zeros([1])])
|
|
|
|
|
|
def get_weightings(weight_schedule, snrs, sigma_data):
|
|
if weight_schedule == "snr":
|
|
weightings = snrs
|
|
elif weight_schedule == "snr+1":
|
|
weightings = snrs + 1
|
|
elif weight_schedule == "karras":
|
|
weightings = snrs + 1.0 / sigma_data**2
|
|
elif weight_schedule == "truncated-snr":
|
|
weightings = th.clamp(snrs, min=1.0)
|
|
elif weight_schedule == "uniform":
|
|
weightings = th.ones_like(snrs)
|
|
else:
|
|
raise NotImplementedError()
|
|
return weightings
|
|
|
|
|
|
class KarrasDenoiser:
|
|
def __init__(
|
|
self,
|
|
sigma_data: float = 0.5,
|
|
sigma_max=80.0,
|
|
sigma_min=0.002,
|
|
rho=7.0,
|
|
weight_schedule="karras",
|
|
distillation=False,
|
|
loss_norm="l2",
|
|
):
|
|
self.sigma_data = sigma_data
|
|
self.sigma_max = sigma_max
|
|
self.sigma_min = sigma_min
|
|
self.weight_schedule = weight_schedule
|
|
self.distillation = distillation
|
|
self.loss_norm = loss_norm
|
|
# if loss_norm == "lpips":
|
|
# self.lpips_loss = LPIPS(replace_pooling=True, reduction="none")
|
|
if loss_norm == "ssim":
|
|
self.ssim_loss = SSIM()
|
|
self.rho = rho
|
|
self.num_timesteps = 40
|
|
|
|
def get_snr(self, sigmas):
|
|
return sigmas**-2
|
|
|
|
def get_sigmas(self, sigmas):
|
|
return sigmas
|
|
|
|
def get_scalings(self, sigma):
|
|
c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
|
|
c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5
|
|
c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
|
|
return c_skip, c_out, c_in
|
|
|
|
def get_scalings_for_boundary_condition(self, sigma):
|
|
c_skip = self.sigma_data**2 / (
|
|
(sigma - self.sigma_min) ** 2 + self.sigma_data**2
|
|
)
|
|
c_out = (
|
|
(sigma - self.sigma_min)
|
|
* self.sigma_data
|
|
/ (sigma**2 + self.sigma_data**2) ** 0.5
|
|
)
|
|
c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
|
|
return c_skip, c_out, c_in
|
|
|
|
def training_losses(self, model, x_start, sigmas, condition=None, noise=None):
|
|
if noise is None:
|
|
noise = th.randn_like(x_start)
|
|
|
|
terms = {}
|
|
|
|
dims = x_start.ndim
|
|
x_t = x_start + noise * append_dims(sigmas, dims)
|
|
model_output, denoised = self.denoise(model, x_t, sigmas, condition)
|
|
|
|
snrs = self.get_snr(sigmas)
|
|
weights = append_dims(
|
|
get_weightings(self.weight_schedule, snrs, self.sigma_data), dims
|
|
)
|
|
# terms["xs_mse"] = mean_flat((denoised - x_start) ** 2)
|
|
terms["mse"] = mean_flat(weights * (denoised - x_start) ** 2)
|
|
# terms["mae"] = mean_flat(weights * th.abs(denoised - x_start))
|
|
# terms["mse"] = nn.MSELoss(reduction="none")(denoised, x_start)
|
|
|
|
# if "vb" in terms:
|
|
# terms["loss"] = terms["mse"] + terms["vb"]
|
|
# else:
|
|
terms["loss"] = terms["mse"]
|
|
|
|
return terms
|
|
|
|
def consistency_losses(
|
|
self,
|
|
model,
|
|
x_start,
|
|
num_scales,
|
|
# model_kwargs=None,
|
|
condition=None,
|
|
target_model=None,
|
|
teacher_model=None,
|
|
teacher_diffusion=None,
|
|
noise=None,
|
|
):
|
|
if noise is None:
|
|
noise = th.randn_like(x_start)
|
|
|
|
dims = x_start.ndim
|
|
|
|
def denoise_fn(x, t):
|
|
return self.denoise(model, x, t, condition)[1]
|
|
|
|
if target_model:
|
|
|
|
@th.no_grad()
|
|
def target_denoise_fn(x, t):
|
|
return self.denoise(target_model, x, t, condition)[1]
|
|
|
|
else:
|
|
raise NotImplementedError("Must have a target model")
|
|
|
|
if teacher_model:
|
|
|
|
@th.no_grad()
|
|
def teacher_denoise_fn(x, t):
|
|
return teacher_diffusion.denoise(teacher_model, x, t, condition)[1]
|
|
|
|
@th.no_grad()
|
|
def heun_solver(samples, t, next_t, x0):
|
|
x = samples
|
|
if teacher_model is None:
|
|
denoiser = x0
|
|
else:
|
|
denoiser = teacher_denoise_fn(x, t)
|
|
|
|
d = (x - denoiser) / append_dims(t, dims)
|
|
samples = x + d * append_dims(next_t - t, dims)
|
|
if teacher_model is None:
|
|
denoiser = x0
|
|
else:
|
|
denoiser = teacher_denoise_fn(samples, next_t)
|
|
|
|
next_d = (samples - denoiser) / append_dims(next_t, dims)
|
|
samples = x + (d + next_d) * append_dims((next_t - t) / 2, dims)
|
|
|
|
return samples
|
|
|
|
@th.no_grad()
|
|
def euler_solver(samples, t, next_t, x0):
|
|
x = samples
|
|
if teacher_model is None:
|
|
denoiser = x0
|
|
else:
|
|
denoiser = teacher_denoise_fn(x, t)
|
|
d = (x - denoiser) / append_dims(t, dims)
|
|
samples = x + d * append_dims(next_t - t, dims)
|
|
|
|
return samples
|
|
|
|
indices = th.randint(
|
|
0, num_scales - 1, (x_start.shape[0],), device=x_start.device
|
|
)
|
|
|
|
t = self.sigma_max ** (1 / self.rho) + indices / (num_scales - 1) * (
|
|
self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
|
|
)
|
|
t = t**self.rho
|
|
|
|
t2 = self.sigma_max ** (1 / self.rho) + (indices + 1) / (num_scales - 1) * (
|
|
self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
|
|
)
|
|
t2 = t2**self.rho
|
|
|
|
x_t = x_start + noise * append_dims(t, dims)
|
|
|
|
dropout_state = th.get_rng_state()
|
|
distiller = denoise_fn(x_t, t)
|
|
|
|
if teacher_model is None:
|
|
x_t2 = euler_solver(x_t, t, t2, x_start).detach()
|
|
else:
|
|
x_t2 = heun_solver(x_t, t, t2, x_start).detach()
|
|
|
|
th.set_rng_state(dropout_state)
|
|
distiller_target = target_denoise_fn(x_t2, t2)
|
|
distiller_target = distiller_target.detach()
|
|
|
|
snrs = self.get_snr(t)
|
|
weights = get_weightings(self.weight_schedule, snrs, self.sigma_data)
|
|
if self.loss_norm == "l1":
|
|
diffs = th.abs(distiller - distiller_target)
|
|
loss = mean_flat(diffs) * weights
|
|
elif self.loss_norm == "l2":
|
|
# diffs = (distiller - distiller_target) ** 2
|
|
loss = F.mse_loss(distiller, distiller_target)
|
|
# loss = mean_flat(diffs) * weights
|
|
elif self.loss_norm == "ssim":
|
|
loss = self.ssim_loss(distiller, distiller_target) * weights
|
|
# elif self.loss_norm == "l2-32":
|
|
# distiller = F.interpolate(distiller, size=32, mode="bilinear")
|
|
# distiller_target = F.interpolate(
|
|
# distiller_target,
|
|
# size=32,
|
|
# mode="bilinear",
|
|
# )
|
|
# diffs = (distiller - distiller_target) ** 2
|
|
# loss = mean_flat(diffs) * weights
|
|
# elif self.loss_norm == "lpips":
|
|
# if x_start.shape[-1] < 256:
|
|
# distiller = F.interpolate(distiller, size=224, mode="bilinear")
|
|
# distiller_target = F.interpolate(
|
|
# distiller_target, size=224, mode="bilinear"
|
|
# )
|
|
|
|
# loss = (
|
|
# self.lpips_loss(
|
|
# (distiller + 1) / 2.0,
|
|
# (distiller_target + 1) / 2.0,
|
|
# )
|
|
# * weights
|
|
# )
|
|
else:
|
|
raise ValueError(f"Unknown loss norm {self.loss_norm}")
|
|
|
|
terms = {}
|
|
terms["loss"] = loss
|
|
|
|
return terms
|
|
|
|
# def progdist_losses(
|
|
# self,
|
|
# model,
|
|
# x_start,
|
|
# num_scales,
|
|
# model_kwargs=None,
|
|
# teacher_model=None,
|
|
# teacher_diffusion=None,
|
|
# noise=None,
|
|
# ):
|
|
# if model_kwargs is None:
|
|
# model_kwargs = {}
|
|
# if noise is None:
|
|
# noise = th.randn_like(x_start)
|
|
|
|
# dims = x_start.ndim
|
|
|
|
# def denoise_fn(x, t):
|
|
# return self.denoise(model, x, t, **model_kwargs)[1]
|
|
|
|
# @th.no_grad()
|
|
# def teacher_denoise_fn(x, t):
|
|
# return teacher_diffusion.denoise(teacher_model, x, t, **model_kwargs)[1]
|
|
|
|
# @th.no_grad()
|
|
# def euler_solver(samples, t, next_t):
|
|
# x = samples
|
|
# denoiser = teacher_denoise_fn(x, t)
|
|
# d = (x - denoiser) / append_dims(t, dims)
|
|
# samples = x + d * append_dims(next_t - t, dims)
|
|
|
|
# return samples
|
|
|
|
# @th.no_grad()
|
|
# def euler_to_denoiser(x_t, t, x_next_t, next_t):
|
|
# denoiser = x_t - append_dims(t, dims) * (x_next_t - x_t) / append_dims(
|
|
# next_t - t, dims
|
|
# )
|
|
# return denoiser
|
|
|
|
# indices = th.randint(0, num_scales, (x_start.shape[0],), device=x_start.device)
|
|
|
|
# t = self.sigma_max ** (1 / self.rho) + indices / num_scales * (
|
|
# self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
|
|
# )
|
|
# t = t**self.rho
|
|
|
|
# t2 = self.sigma_max ** (1 / self.rho) + (indices + 0.5) / num_scales * (
|
|
# self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
|
|
# )
|
|
# t2 = t2**self.rho
|
|
|
|
# t3 = self.sigma_max ** (1 / self.rho) + (indices + 1) / num_scales * (
|
|
# self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
|
|
# )
|
|
# t3 = t3**self.rho
|
|
|
|
# x_t = x_start + noise * append_dims(t, dims)
|
|
|
|
# denoised_x = denoise_fn(x_t, t)
|
|
|
|
# x_t2 = euler_solver(x_t, t, t2).detach()
|
|
# x_t3 = euler_solver(x_t2, t2, t3).detach()
|
|
|
|
# target_x = euler_to_denoiser(x_t, t, x_t3, t3).detach()
|
|
|
|
# snrs = self.get_snr(t)
|
|
# weights = get_weightings(self.weight_schedule, snrs, self.sigma_data)
|
|
# if self.loss_norm == "l1":
|
|
# diffs = th.abs(denoised_x - target_x)
|
|
# loss = mean_flat(diffs) * weights
|
|
# elif self.loss_norm == "l2":
|
|
# diffs = (denoised_x - target_x) ** 2
|
|
# loss = mean_flat(diffs) * weights
|
|
# elif self.loss_norm == "lpips":
|
|
# if x_start.shape[-1] < 256:
|
|
# denoised_x = F.interpolate(denoised_x, size=224, mode="bilinear")
|
|
# target_x = F.interpolate(target_x, size=224, mode="bilinear")
|
|
# loss = (
|
|
# self.lpips_loss(
|
|
# (denoised_x + 1) / 2.0,
|
|
# (target_x + 1) / 2.0,
|
|
# )
|
|
# * weights
|
|
# )
|
|
# else:
|
|
# raise ValueError(f"Unknown loss norm {self.loss_norm}")
|
|
|
|
# terms = {}
|
|
# terms["loss"] = loss
|
|
|
|
# return terms
|
|
|
|
def denoise(self, model, x_t, sigmas, condition):
|
|
if not self.distillation:
|
|
c_skip, c_out, c_in = [
|
|
append_dims(x, x_t.ndim) for x in self.get_scalings(sigmas)
|
|
]
|
|
else:
|
|
c_skip, c_out, c_in = [
|
|
append_dims(x, x_t.ndim)
|
|
for x in self.get_scalings_for_boundary_condition(sigmas)
|
|
]
|
|
rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
|
|
# rescaled_t = rescaled_t[:, None]
|
|
model_output = model(c_in * x_t, rescaled_t, condition)
|
|
denoised = c_out * model_output + c_skip * x_t
|
|
return model_output, denoised
|
|
|
|
|
|
def karras_sample(
|
|
diffusion,
|
|
model,
|
|
shape,
|
|
steps,
|
|
clip_denoised=True,
|
|
progress=True,
|
|
callback=None,
|
|
# model_kwargs=None,
|
|
condition=None,
|
|
device=None,
|
|
sigma_min=0.002,
|
|
sigma_max=80, # higher for highres?
|
|
rho=7.0,
|
|
sampler="heun",
|
|
s_churn=0.0,
|
|
s_tmin=0.0,
|
|
s_tmax=float("inf"),
|
|
s_noise=1.0,
|
|
generator=None,
|
|
ts=None,
|
|
):
|
|
if generator is None:
|
|
generator = get_generator("dummy")
|
|
|
|
if sampler == "progdist":
|
|
sigmas = get_sigmas_karras(steps + 1, sigma_min, sigma_max, rho, device=device)
|
|
else:
|
|
sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
|
|
th.manual_seed(42)
|
|
x_T = generator.randn(*shape, device=device) * sigma_max
|
|
sigmas = sigmas.unsqueeze(-1)
|
|
sample_fn = {
|
|
"heun": sample_heun,
|
|
"dpm": sample_dpm,
|
|
"ancestral": sample_euler_ancestral,
|
|
"onestep": sample_onestep,
|
|
"progdist": sample_progdist,
|
|
"euler": sample_euler,
|
|
"multistep": stochastic_iterative_sampler,
|
|
}[sampler]
|
|
|
|
if sampler in ["heun", "dpm"]:
|
|
sampler_args = dict(
|
|
s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise
|
|
)
|
|
elif sampler == "multistep":
|
|
sampler_args = dict(
|
|
ts=ts, t_min=sigma_min, t_max=sigma_max, rho=diffusion.rho, steps=steps
|
|
)
|
|
else:
|
|
sampler_args = {}
|
|
|
|
def denoiser(x_t, sigma):
|
|
_, denoised = diffusion.denoise(model, x_t, sigma, condition)
|
|
if clip_denoised:
|
|
denoised = denoised.clamp(-1, 1)
|
|
return denoised
|
|
|
|
x_0 = sample_fn(
|
|
denoiser,
|
|
x_T,
|
|
sigmas,
|
|
generator,
|
|
progress=progress,
|
|
callback=callback,
|
|
**sampler_args,
|
|
)
|
|
return x_0.clamp(-1, 1)
|
|
|
|
|
|
def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu"):
|
|
"""Constructs the noise schedule of Karras et al. (2022)."""
|
|
ramp = th.linspace(0, 1, n)
|
|
min_inv_rho = sigma_min ** (1 / rho)
|
|
max_inv_rho = sigma_max ** (1 / rho)
|
|
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
|
|
return append_zero(sigmas).to(device)
|
|
|
|
|
|
def to_d(x, sigma, denoised):
|
|
"""Converts a denoiser output to a Karras ODE derivative."""
|
|
return (x - denoised) / append_dims(sigma, x.ndim)
|
|
|
|
|
|
def get_ancestral_step(sigma_from, sigma_to):
|
|
"""Calculates the noise level (sigma_down) to step down to and the amount
|
|
of noise to add (sigma_up) when doing an ancestral sampling step."""
|
|
sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
|
|
sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
|
|
return sigma_down, sigma_up
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_euler_ancestral(model, x, sigmas, generator, progress=False, callback=None):
|
|
"""Ancestral sampling with Euler method steps."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
indices = range(len(sigmas) - 1)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
indices = tqdm(indices)
|
|
|
|
for i in indices:
|
|
denoised = model(x, sigmas[i] * s_in)
|
|
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
|
|
if callback is not None:
|
|
callback(
|
|
{
|
|
"x": x,
|
|
"i": i,
|
|
"sigma": sigmas[i],
|
|
"sigma_hat": sigmas[i],
|
|
"denoised": denoised,
|
|
}
|
|
)
|
|
d = to_d(x, sigmas[i], denoised)
|
|
# Euler method
|
|
dt = sigma_down - sigmas[i]
|
|
x = x + d * dt
|
|
x = x + generator.randn_like(x) * sigma_up
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_midpoint_ancestral(model, x, ts, generator, progress=False, callback=None):
|
|
"""Ancestral sampling with midpoint method steps."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
step_size = 1 / len(ts)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
ts = tqdm(ts)
|
|
|
|
for tn in ts:
|
|
dn = model(x, tn * s_in)
|
|
dn_2 = model(x + (step_size / 2) * dn, (tn + step_size / 2) * s_in)
|
|
x = x + step_size * dn_2
|
|
if callback is not None:
|
|
callback({"x": x, "tn": tn, "dn": dn, "dn_2": dn_2})
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_heun(
|
|
denoiser,
|
|
x,
|
|
sigmas,
|
|
generator,
|
|
progress=False,
|
|
callback=None,
|
|
s_churn=0.0,
|
|
s_tmin=0.0,
|
|
s_tmax=float("inf"),
|
|
s_noise=1.0,
|
|
):
|
|
"""Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
indices = range(len(sigmas) - 1)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
indices = tqdm(indices)
|
|
|
|
for i in indices:
|
|
gamma = (
|
|
min(s_churn / (len(sigmas) - 1), 2**0.5 - 1)
|
|
if s_tmin <= sigmas[i] <= s_tmax
|
|
else 0.0
|
|
)
|
|
eps = generator.randn_like(x) * s_noise
|
|
sigma_hat = sigmas[i] * (gamma + 1)
|
|
if gamma > 0:
|
|
x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
|
|
denoised = denoiser(x, sigma_hat * s_in)
|
|
d = to_d(x, sigma_hat, denoised)
|
|
if callback is not None:
|
|
callback(
|
|
{
|
|
"x": x,
|
|
"i": i,
|
|
"sigma": sigmas[i],
|
|
"sigma_hat": sigma_hat,
|
|
"denoised": denoised,
|
|
}
|
|
)
|
|
dt = sigmas[i + 1] - sigma_hat
|
|
if sigmas[i + 1] == 0:
|
|
# Euler method
|
|
x = x + d * dt
|
|
else:
|
|
# Heun's method
|
|
x_2 = x + d * dt
|
|
denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
|
|
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
|
|
d_prime = (d + d_2) / 2
|
|
x = x + d_prime * dt
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_euler(
|
|
denoiser,
|
|
x,
|
|
sigmas,
|
|
generator,
|
|
progress=False,
|
|
callback=None,
|
|
):
|
|
"""Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
indices = range(len(sigmas) - 1)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
indices = tqdm(indices)
|
|
|
|
for i in indices:
|
|
sigma = sigmas[i]
|
|
denoised = denoiser(x, sigma * s_in)
|
|
d = to_d(x, sigma, denoised)
|
|
if callback is not None:
|
|
callback(
|
|
{
|
|
"x": x,
|
|
"i": i,
|
|
"sigma": sigmas[i],
|
|
"denoised": denoised,
|
|
}
|
|
)
|
|
dt = sigmas[i + 1] - sigma
|
|
x = x + d * dt
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_dpm(
|
|
denoiser,
|
|
x,
|
|
sigmas,
|
|
generator,
|
|
progress=False,
|
|
callback=None,
|
|
s_churn=0.0,
|
|
s_tmin=0.0,
|
|
s_tmax=float("inf"),
|
|
s_noise=1.0,
|
|
):
|
|
"""A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
indices = range(len(sigmas) - 1)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
indices = tqdm(indices)
|
|
|
|
for i in indices:
|
|
gamma = (
|
|
min(s_churn / (len(sigmas) - 1), 2**0.5 - 1)
|
|
if s_tmin <= sigmas[i] <= s_tmax
|
|
else 0.0
|
|
)
|
|
eps = generator.randn_like(x) * s_noise
|
|
sigma_hat = sigmas[i] * (gamma + 1)
|
|
if gamma > 0:
|
|
x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
|
|
denoised = denoiser(x, sigma_hat * s_in)
|
|
d = to_d(x, sigma_hat, denoised)
|
|
if callback is not None:
|
|
callback(
|
|
{
|
|
"x": x,
|
|
"i": i,
|
|
"sigma": sigmas[i],
|
|
"sigma_hat": sigma_hat,
|
|
"denoised": denoised,
|
|
}
|
|
)
|
|
# Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
|
|
sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2) ** 3
|
|
dt_1 = sigma_mid - sigma_hat
|
|
dt_2 = sigmas[i + 1] - sigma_hat
|
|
x_2 = x + d * dt_1
|
|
denoised_2 = denoiser(x_2, sigma_mid * s_in)
|
|
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
|
x = x + d_2 * dt_2
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_onestep(
|
|
distiller,
|
|
x,
|
|
sigmas,
|
|
generator=None,
|
|
progress=False,
|
|
callback=None,
|
|
):
|
|
"""Single-step generation from a distilled model."""
|
|
s_in = x.new_ones([x.shape[0]])
|
|
return distiller(x, sigmas[0] * s_in)
|
|
|
|
|
|
@th.no_grad()
|
|
def stochastic_iterative_sampler(
|
|
distiller,
|
|
x,
|
|
sigmas,
|
|
generator,
|
|
ts,
|
|
progress=False,
|
|
callback=None,
|
|
t_min=0.002,
|
|
t_max=80.0,
|
|
rho=7.0,
|
|
steps=40,
|
|
):
|
|
t_max_rho = t_max ** (1 / rho)
|
|
t_min_rho = t_min ** (1 / rho)
|
|
s_in = x.new_ones([x.shape[0]])
|
|
|
|
for i in range(len(ts) - 1):
|
|
t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
x0 = distiller(x, t * s_in)
|
|
next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
next_t = np.clip(next_t, t_min, t_max)
|
|
x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)
|
|
|
|
return x
|
|
|
|
|
|
@th.no_grad()
|
|
def sample_progdist(
|
|
denoiser,
|
|
x,
|
|
sigmas,
|
|
generator=None,
|
|
progress=False,
|
|
callback=None,
|
|
):
|
|
s_in = x.new_ones([x.shape[0]])
|
|
sigmas = sigmas[:-1] # skip the zero sigma
|
|
|
|
indices = range(len(sigmas) - 1)
|
|
if progress:
|
|
from tqdm.auto import tqdm
|
|
|
|
indices = tqdm(indices)
|
|
|
|
for i in indices:
|
|
sigma = sigmas[i]
|
|
denoised = denoiser(x, sigma * s_in)
|
|
d = to_d(x, sigma, denoised)
|
|
if callback is not None:
|
|
callback(
|
|
{
|
|
"x": x,
|
|
"i": i,
|
|
"sigma": sigma,
|
|
"denoised": denoised,
|
|
}
|
|
)
|
|
dt = sigmas[i + 1] - sigma
|
|
x = x + d * dt
|
|
|
|
return x
|
|
|
|
|
|
# @th.no_grad()
|
|
# def iterative_colorization(
|
|
# distiller,
|
|
# images,
|
|
# x,
|
|
# ts,
|
|
# t_min=0.002,
|
|
# t_max=80.0,
|
|
# rho=7.0,
|
|
# steps=40,
|
|
# generator=None,
|
|
# ):
|
|
# def obtain_orthogonal_matrix():
|
|
# vector = np.asarray([0.2989, 0.5870, 0.1140])
|
|
# vector = vector / np.linalg.norm(vector)
|
|
# matrix = np.eye(3)
|
|
# matrix[:, 0] = vector
|
|
# matrix = np.linalg.qr(matrix)[0]
|
|
# if np.sum(matrix[:, 0]) < 0:
|
|
# matrix = -matrix
|
|
# return matrix
|
|
|
|
# Q = th.from_numpy(obtain_orthogonal_matrix()).to(dist_util.dev()).to(th.float32)
|
|
# mask = th.zeros(*x.shape[1:], device=dist_util.dev())
|
|
# mask[0, ...] = 1.0
|
|
|
|
# def replacement(x0, x1):
|
|
# x0 = th.einsum("bchw,cd->bdhw", x0, Q)
|
|
# x1 = th.einsum("bchw,cd->bdhw", x1, Q)
|
|
|
|
# x_mix = x0 * mask + x1 * (1.0 - mask)
|
|
# x_mix = th.einsum("bdhw,cd->bchw", x_mix, Q)
|
|
# return x_mix
|
|
|
|
# t_max_rho = t_max ** (1 / rho)
|
|
# t_min_rho = t_min ** (1 / rho)
|
|
# s_in = x.new_ones([x.shape[0]])
|
|
# images = replacement(images, th.zeros_like(images))
|
|
|
|
# for i in range(len(ts) - 1):
|
|
# t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# x0 = distiller(x, t * s_in)
|
|
# x0 = th.clamp(x0, -1.0, 1.0)
|
|
# x0 = replacement(images, x0)
|
|
# next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# next_t = np.clip(next_t, t_min, t_max)
|
|
# x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)
|
|
|
|
# return x, images
|
|
|
|
|
|
# @th.no_grad()
|
|
# def iterative_inpainting(
|
|
# distiller,
|
|
# images,
|
|
# x,
|
|
# ts,
|
|
# t_min=0.002,
|
|
# t_max=80.0,
|
|
# rho=7.0,
|
|
# steps=40,
|
|
# generator=None,
|
|
# ):
|
|
# from PIL import Image, ImageDraw, ImageFont
|
|
|
|
# image_size = x.shape[-1]
|
|
|
|
# # create a blank image with a white background
|
|
# img = Image.new("RGB", (image_size, image_size), color="white")
|
|
|
|
# # get a drawing context for the image
|
|
# draw = ImageDraw.Draw(img)
|
|
|
|
# # load a font
|
|
# font = ImageFont.truetype("arial.ttf", 250)
|
|
|
|
# # draw the letter "C" in black
|
|
# draw.text((50, 0), "S", font=font, fill=(0, 0, 0))
|
|
|
|
# # convert the image to a numpy array
|
|
# img_np = np.array(img)
|
|
# img_np = img_np.transpose(2, 0, 1)
|
|
# img_th = th.from_numpy(img_np).to(dist_util.dev())
|
|
|
|
# mask = th.zeros(*x.shape, device=dist_util.dev())
|
|
# mask = mask.reshape(-1, 7, 3, image_size, image_size)
|
|
|
|
# mask[::2, :, img_th > 0.5] = 1.0
|
|
# mask[1::2, :, img_th < 0.5] = 1.0
|
|
# mask = mask.reshape(-1, 3, image_size, image_size)
|
|
|
|
# def replacement(x0, x1):
|
|
# x_mix = x0 * mask + x1 * (1 - mask)
|
|
# return x_mix
|
|
|
|
# t_max_rho = t_max ** (1 / rho)
|
|
# t_min_rho = t_min ** (1 / rho)
|
|
# s_in = x.new_ones([x.shape[0]])
|
|
# images = replacement(images, -th.ones_like(images))
|
|
|
|
# for i in range(len(ts) - 1):
|
|
# t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# x0 = distiller(x, t * s_in)
|
|
# x0 = th.clamp(x0, -1.0, 1.0)
|
|
# x0 = replacement(images, x0)
|
|
# next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# next_t = np.clip(next_t, t_min, t_max)
|
|
# x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)
|
|
|
|
# return x, images
|
|
|
|
|
|
# @th.no_grad()
|
|
# def iterative_superres(
|
|
# distiller,
|
|
# images,
|
|
# x,
|
|
# ts,
|
|
# t_min=0.002,
|
|
# t_max=80.0,
|
|
# rho=7.0,
|
|
# steps=40,
|
|
# generator=None,
|
|
# ):
|
|
# patch_size = 8
|
|
|
|
# def obtain_orthogonal_matrix():
|
|
# vector = np.asarray([1] * patch_size**2)
|
|
# vector = vector / np.linalg.norm(vector)
|
|
# matrix = np.eye(patch_size**2)
|
|
# matrix[:, 0] = vector
|
|
# matrix = np.linalg.qr(matrix)[0]
|
|
# if np.sum(matrix[:, 0]) < 0:
|
|
# matrix = -matrix
|
|
# return matrix
|
|
|
|
# Q = th.from_numpy(obtain_orthogonal_matrix()).to(dist_util.dev()).to(th.float32)
|
|
|
|
# image_size = x.shape[-1]
|
|
|
|
# def replacement(x0, x1):
|
|
# x0_flatten = (
|
|
# x0.reshape(-1, 3, image_size, image_size)
|
|
# .reshape(
|
|
# -1,
|
|
# 3,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# )
|
|
# .permute(0, 1, 2, 4, 3, 5)
|
|
# .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
|
|
# )
|
|
# x1_flatten = (
|
|
# x1.reshape(-1, 3, image_size, image_size)
|
|
# .reshape(
|
|
# -1,
|
|
# 3,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# )
|
|
# .permute(0, 1, 2, 4, 3, 5)
|
|
# .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
|
|
# )
|
|
# x0 = th.einsum("bcnd,de->bcne", x0_flatten, Q)
|
|
# x1 = th.einsum("bcnd,de->bcne", x1_flatten, Q)
|
|
# x_mix = x0.new_zeros(x0.shape)
|
|
# x_mix[..., 0] = x0[..., 0]
|
|
# x_mix[..., 1:] = x1[..., 1:]
|
|
# x_mix = th.einsum("bcne,de->bcnd", x_mix, Q)
|
|
# x_mix = (
|
|
# x_mix.reshape(
|
|
# -1,
|
|
# 3,
|
|
# image_size // patch_size,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# patch_size,
|
|
# )
|
|
# .permute(0, 1, 2, 4, 3, 5)
|
|
# .reshape(-1, 3, image_size, image_size)
|
|
# )
|
|
# return x_mix
|
|
|
|
# def average_image_patches(x):
|
|
# x_flatten = (
|
|
# x.reshape(-1, 3, image_size, image_size)
|
|
# .reshape(
|
|
# -1,
|
|
# 3,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# )
|
|
# .permute(0, 1, 2, 4, 3, 5)
|
|
# .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
|
|
# )
|
|
# x_flatten[..., :] = x_flatten.mean(dim=-1, keepdim=True)
|
|
# return (
|
|
# x_flatten.reshape(
|
|
# -1,
|
|
# 3,
|
|
# image_size // patch_size,
|
|
# image_size // patch_size,
|
|
# patch_size,
|
|
# patch_size,
|
|
# )
|
|
# .permute(0, 1, 2, 4, 3, 5)
|
|
# .reshape(-1, 3, image_size, image_size)
|
|
# )
|
|
|
|
# t_max_rho = t_max ** (1 / rho)
|
|
# t_min_rho = t_min ** (1 / rho)
|
|
# s_in = x.new_ones([x.shape[0]])
|
|
# images = average_image_patches(images)
|
|
|
|
# for i in range(len(ts) - 1):
|
|
# t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# x0 = distiller(x, t * s_in)
|
|
# x0 = th.clamp(x0, -1.0, 1.0)
|
|
# x0 = replacement(images, x0)
|
|
# next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
|
|
# next_t = np.clip(next_t, t_min, t_max)
|
|
# x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)
|
|
|
|
# return x, images
|