450 lines
15 KiB
Python
450 lines
15 KiB
Python
|
"""
|
||
|
|
Partially ported from https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
from typing import Dict, Union
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from omegaconf import ListConfig, OmegaConf
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from ...modules.diffusionmodules.sampling_utils import (
|
||
|
|
get_ancestral_step,
|
||
|
|
linear_multistep_coeff,
|
||
|
|
to_d,
|
||
|
|
to_neg_log_sigma,
|
||
|
|
to_sigma,
|
||
|
|
)
|
||
|
|
from ...util import append_dims, default, instantiate_from_config
|
||
|
|
|
||
|
|
DEFAULT_GUIDER = {"target": "sgm.modules.diffusionmodules.guiders.IdentityGuider"}
|
||
|
|
|
||
|
|
|
||
|
|
class BaseDiffusionSampler:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
discretization_config: Union[Dict, ListConfig, OmegaConf],
|
||
|
|
num_steps: Union[int, None] = None,
|
||
|
|
guider_config: Union[Dict, ListConfig, OmegaConf, None] = None,
|
||
|
|
verbose: bool = False,
|
||
|
|
device: str = "cuda",
|
||
|
|
):
|
||
|
|
self.num_steps = num_steps
|
||
|
|
self.discretization = instantiate_from_config(discretization_config)
|
||
|
|
self.guider = instantiate_from_config(
|
||
|
|
default(
|
||
|
|
guider_config,
|
||
|
|
DEFAULT_GUIDER,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self.verbose = verbose
|
||
|
|
self.device = device
|
||
|
|
|
||
|
|
def prepare_sampling_loop(self, x, cond, uc=None, num_steps=None):
|
||
|
|
sigmas = self.discretization(
|
||
|
|
self.num_steps if num_steps is None else num_steps, device=self.device
|
||
|
|
)
|
||
|
|
uc = default(uc, cond)
|
||
|
|
|
||
|
|
x *= torch.sqrt(1.0 + sigmas[0] ** 2.0)
|
||
|
|
num_sigmas = len(sigmas)
|
||
|
|
|
||
|
|
s_in = x.new_ones([x.shape[0]])
|
||
|
|
|
||
|
|
return x, s_in, sigmas, num_sigmas, cond, uc
|
||
|
|
|
||
|
|
def denoise(self, x, denoiser, sigma, cond, uc):
|
||
|
|
denoised = denoiser(*self.guider.prepare_inputs(x, sigma, cond, uc))
|
||
|
|
denoised = self.guider(denoised, sigma)
|
||
|
|
return denoised
|
||
|
|
|
||
|
|
def get_sigma_gen(self, num_sigmas):
|
||
|
|
sigma_generator = range(num_sigmas - 1)
|
||
|
|
if self.verbose:
|
||
|
|
print("#" * 30, " Sampling setting ", "#" * 30)
|
||
|
|
print(f"Sampler: {self.__class__.__name__}")
|
||
|
|
print(f"Discretization: {self.discretization.__class__.__name__}")
|
||
|
|
print(f"Guider: {self.guider.__class__.__name__}")
|
||
|
|
sigma_generator = tqdm(
|
||
|
|
sigma_generator,
|
||
|
|
total=num_sigmas,
|
||
|
|
desc=f"Sampling with {self.__class__.__name__} for {num_sigmas} steps",
|
||
|
|
)
|
||
|
|
return sigma_generator
|
||
|
|
|
||
|
|
|
||
|
|
class SingleStepDiffusionSampler(BaseDiffusionSampler):
|
||
|
|
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc, *args, **kwargs):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def euler_step(self, x, d, dt):
|
||
|
|
return x + dt * d
|
||
|
|
|
||
|
|
|
||
|
|
class EDMSampler(SingleStepDiffusionSampler):
|
||
|
|
def __init__(
|
||
|
|
self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, *args, **kwargs
|
||
|
|
):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self.s_churn = s_churn
|
||
|
|
self.s_tmin = s_tmin
|
||
|
|
self.s_tmax = s_tmax
|
||
|
|
self.s_noise = s_noise
|
||
|
|
|
||
|
|
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, gamma=0.0):
|
||
|
|
sigma_hat = sigma * (gamma + 1.0)
|
||
|
|
if gamma > 0:
|
||
|
|
eps = torch.randn_like(x) * self.s_noise
|
||
|
|
x = x + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5
|
||
|
|
|
||
|
|
denoised = self.denoise(x, denoiser, sigma_hat, cond, uc)
|
||
|
|
# print('denoised', denoised.mean(axis=[0, 2, 3]))
|
||
|
|
d = to_d(x, sigma_hat, denoised)
|
||
|
|
dt = append_dims(next_sigma - sigma_hat, x.ndim)
|
||
|
|
|
||
|
|
euler_step = self.euler_step(x, d, dt)
|
||
|
|
x = self.possible_correction_step(
|
||
|
|
euler_step, x, d, dt, next_sigma, denoiser, cond, uc
|
||
|
|
)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def __call__(self, denoiser, x, cond, uc=None, num_steps=None):
|
||
|
|
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
|
||
|
|
x, cond, uc, num_steps
|
||
|
|
)
|
||
|
|
|
||
|
|
for i in self.get_sigma_gen(num_sigmas):
|
||
|
|
gamma = (
|
||
|
|
min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1)
|
||
|
|
if self.s_tmin <= sigmas[i] <= self.s_tmax
|
||
|
|
else 0.0
|
||
|
|
)
|
||
|
|
x = self.sampler_step(
|
||
|
|
s_in * sigmas[i],
|
||
|
|
s_in * sigmas[i + 1],
|
||
|
|
denoiser,
|
||
|
|
x,
|
||
|
|
cond,
|
||
|
|
uc,
|
||
|
|
gamma,
|
||
|
|
)
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class AncestralSampler(SingleStepDiffusionSampler):
|
||
|
|
def __init__(self, eta=1.0, s_noise=1.0, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self.eta = eta
|
||
|
|
self.s_noise = s_noise
|
||
|
|
self.noise_sampler = lambda x: torch.randn_like(x)
|
||
|
|
|
||
|
|
def ancestral_euler_step(self, x, denoised, sigma, sigma_down):
|
||
|
|
d = to_d(x, sigma, denoised)
|
||
|
|
dt = append_dims(sigma_down - sigma, x.ndim)
|
||
|
|
|
||
|
|
return self.euler_step(x, d, dt)
|
||
|
|
|
||
|
|
def ancestral_step(self, x, sigma, next_sigma, sigma_up):
|
||
|
|
x = torch.where(
|
||
|
|
append_dims(next_sigma, x.ndim) > 0.0,
|
||
|
|
x + self.noise_sampler(x) * self.s_noise * append_dims(sigma_up, x.ndim),
|
||
|
|
x,
|
||
|
|
)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def __call__(self, denoiser, x, cond, uc=None, num_steps=None):
|
||
|
|
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
|
||
|
|
x, cond, uc, num_steps
|
||
|
|
)
|
||
|
|
|
||
|
|
for i in self.get_sigma_gen(num_sigmas):
|
||
|
|
x = self.sampler_step(
|
||
|
|
s_in * sigmas[i],
|
||
|
|
s_in * sigmas[i + 1],
|
||
|
|
denoiser,
|
||
|
|
x,
|
||
|
|
cond,
|
||
|
|
uc,
|
||
|
|
)
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class LinearMultistepSampler(BaseDiffusionSampler):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
order=4,
|
||
|
|
*args,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self.order = order
|
||
|
|
|
||
|
|
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, **kwargs):
|
||
|
|
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
|
||
|
|
x, cond, uc, num_steps
|
||
|
|
)
|
||
|
|
|
||
|
|
ds = []
|
||
|
|
sigmas_cpu = sigmas.detach().cpu().numpy()
|
||
|
|
for i in self.get_sigma_gen(num_sigmas):
|
||
|
|
sigma = s_in * sigmas[i]
|
||
|
|
denoised = denoiser(
|
||
|
|
*self.guider.prepare_inputs(x, sigma, cond, uc), **kwargs
|
||
|
|
)
|
||
|
|
denoised = self.guider(denoised, sigma)
|
||
|
|
d = to_d(x, sigma, denoised)
|
||
|
|
ds.append(d)
|
||
|
|
if len(ds) > self.order:
|
||
|
|
ds.pop(0)
|
||
|
|
cur_order = min(i + 1, self.order)
|
||
|
|
coeffs = [
|
||
|
|
linear_multistep_coeff(cur_order, sigmas_cpu, i, j)
|
||
|
|
for j in range(cur_order)
|
||
|
|
]
|
||
|
|
x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class EulerEDMSampler(EDMSampler):
|
||
|
|
def possible_correction_step(
|
||
|
|
self, euler_step, x, d, dt, next_sigma, denoiser, cond, uc
|
||
|
|
):
|
||
|
|
# print("euler_step: ", euler_step.mean(axis=[0, 2, 3]))
|
||
|
|
return euler_step
|
||
|
|
|
||
|
|
|
||
|
|
class HeunEDMSampler(EDMSampler):
|
||
|
|
def possible_correction_step(
|
||
|
|
self, euler_step, x, d, dt, next_sigma, denoiser, cond, uc
|
||
|
|
):
|
||
|
|
if torch.sum(next_sigma) < 1e-14:
|
||
|
|
# Save a network evaluation if all noise levels are 0
|
||
|
|
return euler_step
|
||
|
|
else:
|
||
|
|
denoised = self.denoise(euler_step, denoiser, next_sigma, cond, uc)
|
||
|
|
d_new = to_d(euler_step, next_sigma, denoised)
|
||
|
|
d_prime = (d + d_new) / 2.0
|
||
|
|
|
||
|
|
# apply correction if noise level is not 0
|
||
|
|
x = torch.where(
|
||
|
|
append_dims(next_sigma, x.ndim) > 0.0, x + d_prime * dt, euler_step
|
||
|
|
)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class EulerAncestralSampler(AncestralSampler):
|
||
|
|
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc):
|
||
|
|
sigma_down, sigma_up = get_ancestral_step(sigma, next_sigma, eta=self.eta)
|
||
|
|
denoised = self.denoise(x, denoiser, sigma, cond, uc)
|
||
|
|
x = self.ancestral_euler_step(x, denoised, sigma, sigma_down)
|
||
|
|
x = self.ancestral_step(x, sigma, next_sigma, sigma_up)
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class DPMPP2SAncestralSampler(AncestralSampler):
|
||
|
|
def get_variables(self, sigma, sigma_down):
|
||
|
|
t, t_next = [to_neg_log_sigma(s) for s in (sigma, sigma_down)]
|
||
|
|
h = t_next - t
|
||
|
|
s = t + 0.5 * h
|
||
|
|
return h, s, t, t_next
|
||
|
|
|
||
|
|
def get_mult(self, h, s, t, t_next):
|
||
|
|
mult1 = to_sigma(s) / to_sigma(t)
|
||
|
|
mult2 = (-0.5 * h).expm1()
|
||
|
|
mult3 = to_sigma(t_next) / to_sigma(t)
|
||
|
|
mult4 = (-h).expm1()
|
||
|
|
|
||
|
|
return mult1, mult2, mult3, mult4
|
||
|
|
|
||
|
|
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, **kwargs):
|
||
|
|
sigma_down, sigma_up = get_ancestral_step(sigma, next_sigma, eta=self.eta)
|
||
|
|
denoised = self.denoise(x, denoiser, sigma, cond, uc)
|
||
|
|
x_euler = self.ancestral_euler_step(x, denoised, sigma, sigma_down)
|
||
|
|
|
||
|
|
if torch.sum(sigma_down) < 1e-14:
|
||
|
|
# Save a network evaluation if all noise levels are 0
|
||
|
|
x = x_euler
|
||
|
|
else:
|
||
|
|
h, s, t, t_next = self.get_variables(sigma, sigma_down)
|
||
|
|
mult = [
|
||
|
|
append_dims(mult, x.ndim) for mult in self.get_mult(h, s, t, t_next)
|
||
|
|
]
|
||
|
|
|
||
|
|
x2 = mult[0] * x - mult[1] * denoised
|
||
|
|
denoised2 = self.denoise(x2, denoiser, to_sigma(s), cond, uc)
|
||
|
|
x_dpmpp2s = mult[2] * x - mult[3] * denoised2
|
||
|
|
|
||
|
|
# apply correction if noise level is not 0
|
||
|
|
x = torch.where(append_dims(sigma_down, x.ndim) > 0.0, x_dpmpp2s, x_euler)
|
||
|
|
|
||
|
|
x = self.ancestral_step(x, sigma, next_sigma, sigma_up)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class DPMPP2MSampler(BaseDiffusionSampler):
|
||
|
|
def get_variables(self, sigma, next_sigma, previous_sigma=None):
|
||
|
|
t, t_next = [to_neg_log_sigma(s) for s in (sigma, next_sigma)]
|
||
|
|
h = t_next - t
|
||
|
|
|
||
|
|
if previous_sigma is not None:
|
||
|
|
h_last = t - to_neg_log_sigma(previous_sigma)
|
||
|
|
r = h_last / h
|
||
|
|
return h, r, t, t_next
|
||
|
|
else:
|
||
|
|
return h, None, t, t_next
|
||
|
|
|
||
|
|
def get_mult(self, h, r, t, t_next, previous_sigma):
|
||
|
|
mult1 = to_sigma(t_next) / to_sigma(t)
|
||
|
|
mult2 = (-h).expm1()
|
||
|
|
|
||
|
|
if previous_sigma is not None:
|
||
|
|
mult3 = 1 + 1 / (2 * r)
|
||
|
|
mult4 = 1 / (2 * r)
|
||
|
|
return mult1, mult2, mult3, mult4
|
||
|
|
else:
|
||
|
|
return mult1, mult2
|
||
|
|
|
||
|
|
def sampler_step(
|
||
|
|
self,
|
||
|
|
old_denoised,
|
||
|
|
previous_sigma,
|
||
|
|
sigma,
|
||
|
|
next_sigma,
|
||
|
|
denoiser,
|
||
|
|
x,
|
||
|
|
cond,
|
||
|
|
uc=None,
|
||
|
|
):
|
||
|
|
denoised = self.denoise(x, denoiser, sigma, cond, uc)
|
||
|
|
|
||
|
|
h, r, t, t_next = self.get_variables(sigma, next_sigma, previous_sigma)
|
||
|
|
mult = [
|
||
|
|
append_dims(mult, x.ndim)
|
||
|
|
for mult in self.get_mult(h, r, t, t_next, previous_sigma)
|
||
|
|
]
|
||
|
|
|
||
|
|
x_standard = mult[0] * x - mult[1] * denoised
|
||
|
|
if old_denoised is None or torch.sum(next_sigma) < 1e-14:
|
||
|
|
# Save a network evaluation if all noise levels are 0 or on the first step
|
||
|
|
return x_standard, denoised
|
||
|
|
else:
|
||
|
|
denoised_d = mult[2] * denoised - mult[3] * old_denoised
|
||
|
|
x_advanced = mult[0] * x - mult[1] * denoised_d
|
||
|
|
|
||
|
|
# apply correction if noise level is not 0 and not first step
|
||
|
|
x = torch.where(
|
||
|
|
append_dims(next_sigma, x.ndim) > 0.0, x_advanced, x_standard
|
||
|
|
)
|
||
|
|
|
||
|
|
return x, denoised
|
||
|
|
|
||
|
|
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, **kwargs):
|
||
|
|
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
|
||
|
|
x, cond, uc, num_steps
|
||
|
|
)
|
||
|
|
|
||
|
|
old_denoised = None
|
||
|
|
for i in self.get_sigma_gen(num_sigmas):
|
||
|
|
x, old_denoised = self.sampler_step(
|
||
|
|
old_denoised,
|
||
|
|
None if i == 0 else s_in * sigmas[i - 1],
|
||
|
|
s_in * sigmas[i],
|
||
|
|
s_in * sigmas[i + 1],
|
||
|
|
denoiser,
|
||
|
|
x,
|
||
|
|
cond,
|
||
|
|
uc=uc,
|
||
|
|
)
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class RestoreEDMSampler(SingleStepDiffusionSampler):
|
||
|
|
def __init__(
|
||
|
|
self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, restore_cfg=4.0,
|
||
|
|
restore_cfg_s_tmin=0.05, *args, **kwargs
|
||
|
|
):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self.s_churn = s_churn
|
||
|
|
self.s_tmin = s_tmin
|
||
|
|
self.s_tmax = s_tmax
|
||
|
|
self.s_noise = s_noise
|
||
|
|
self.restore_cfg = restore_cfg
|
||
|
|
self.restore_cfg_s_tmin = restore_cfg_s_tmin
|
||
|
|
self.sigma_max = 14.6146
|
||
|
|
|
||
|
|
def denoise(self, x, denoiser, sigma, cond, uc, control_scale=1.0):
|
||
|
|
denoised = denoiser(*self.guider.prepare_inputs(x, sigma, cond, uc), control_scale)
|
||
|
|
denoised = self.guider(denoised, sigma)
|
||
|
|
return denoised
|
||
|
|
|
||
|
|
|
||
|
|
def sampler_step(self, sigma, next_sigma, denoiser, x, cond, uc=None, gamma=0.0, x_center=None, eps_noise=None,
|
||
|
|
control_scale=1.0, use_linear_control_scale=False, control_scale_start=0.0):
|
||
|
|
sigma_hat = sigma * (gamma + 1.0)
|
||
|
|
if gamma > 0:
|
||
|
|
if eps_noise is not None:
|
||
|
|
eps = eps_noise * self.s_noise
|
||
|
|
else:
|
||
|
|
eps = torch.randn_like(x) * self.s_noise
|
||
|
|
x = x + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5
|
||
|
|
|
||
|
|
if use_linear_control_scale:
|
||
|
|
control_scale = (sigma[0].item() / self.sigma_max) * (control_scale_start - control_scale) + control_scale
|
||
|
|
|
||
|
|
denoised = self.denoise(x, denoiser, sigma_hat, cond, uc, control_scale=control_scale)
|
||
|
|
|
||
|
|
if (next_sigma[0] > self.restore_cfg_s_tmin) and (self.restore_cfg > 0):
|
||
|
|
d_center = (denoised - x_center)
|
||
|
|
denoised = denoised - d_center * ((sigma.view(-1, 1, 1, 1) / self.sigma_max) ** self.restore_cfg)
|
||
|
|
|
||
|
|
d = to_d(x, sigma_hat, denoised)
|
||
|
|
dt = append_dims(next_sigma - sigma_hat, x.ndim)
|
||
|
|
x = self.euler_step(x, d, dt)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def __call__(self, denoiser, x, cond, uc=None, num_steps=None, x_center=None, control_scale=1.0,
|
||
|
|
use_linear_control_scale=False, control_scale_start=0.0):
|
||
|
|
x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop(
|
||
|
|
x, cond, uc, num_steps
|
||
|
|
)
|
||
|
|
|
||
|
|
for _idx, i in enumerate(self.get_sigma_gen(num_sigmas)):
|
||
|
|
gamma = (
|
||
|
|
min(self.s_churn / (num_sigmas - 1), 2**0.5 - 1)
|
||
|
|
if self.s_tmin <= sigmas[i] <= self.s_tmax
|
||
|
|
else 0.0
|
||
|
|
)
|
||
|
|
x = self.sampler_step(
|
||
|
|
s_in * sigmas[i],
|
||
|
|
s_in * sigmas[i + 1],
|
||
|
|
denoiser,
|
||
|
|
x,
|
||
|
|
cond,
|
||
|
|
uc,
|
||
|
|
gamma,
|
||
|
|
x_center,
|
||
|
|
control_scale=control_scale,
|
||
|
|
use_linear_control_scale=use_linear_control_scale,
|
||
|
|
control_scale_start=control_scale_start,
|
||
|
|
)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def to_d_center(denoised, x_center, x):
|
||
|
|
b = denoised.shape[0]
|
||
|
|
v_center = (denoised - x_center).view(b, -1)
|
||
|
|
v_denoise = (x - denoised).view(b, -1)
|
||
|
|
d_center = v_center - v_denoise * (v_center * v_denoise).sum(dim=1).view(b, 1) / \
|
||
|
|
(v_denoise * v_denoise).sum(dim=1).view(b, 1)
|
||
|
|
d_center = d_center / d_center.view(x.shape[0], -1).norm(dim=1).view(-1, 1)
|
||
|
|
return d_center.view(denoised.shape)
|