41 lines
1.7 KiB
Python
Executable File
41 lines
1.7 KiB
Python
Executable File
# Only include samplers that are not already in A1111
|
|
|
|
import torch
|
|
|
|
from tqdm import trange
|
|
|
|
|
|
def default_noise_sampler(x):
|
|
return lambda sigma, sigma_next: torch.randn_like(x)
|
|
|
|
|
|
def generic_step_sampler(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, step_function=None):
|
|
extra_args = {} if extra_args is None else extra_args
|
|
noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
|
|
s_in = x.new_ones([x.shape[0]])
|
|
|
|
for i in trange(len(sigmas) - 1, disable=disable):
|
|
denoised = model(x, sigmas[i] * s_in, **extra_args)
|
|
if callback is not None:
|
|
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
|
x = step_function(x / torch.sqrt(1.0 + sigmas[i] ** 2.0), sigmas[i], sigmas[i + 1], (x - denoised) / sigmas[i], noise_sampler)
|
|
if sigmas[i + 1] != 0:
|
|
x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2.0)
|
|
return x
|
|
|
|
|
|
def DDPMSampler_step(x, sigma, sigma_prev, noise, noise_sampler):
|
|
alpha_cumprod = 1 / ((sigma * sigma) + 1)
|
|
alpha_cumprod_prev = 1 / ((sigma_prev * sigma_prev) + 1)
|
|
alpha = (alpha_cumprod / alpha_cumprod_prev)
|
|
|
|
mu = (1.0 / alpha).sqrt() * (x - (1 - alpha) * noise / (1 - alpha_cumprod).sqrt())
|
|
if sigma_prev > 0:
|
|
mu += ((1 - alpha) * (1. - alpha_cumprod_prev) / (1. - alpha_cumprod)).sqrt() * noise_sampler(sigma, sigma_prev)
|
|
return mu
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_ddpm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
|
|
return generic_step_sampler(model, x, sigmas, extra_args, callback, disable, noise_sampler, DDPMSampler_step)
|