155 lines
6.3 KiB
Python
155 lines
6.3 KiB
Python
import torch
|
|
|
|
class ERSDEScheduler():
|
|
"""Extended Reverse-Time SDE solver (VP ER-SDE-Solver-3).
|
|
|
|
Based on: arXiv: https://arxiv.org/abs/2309.06169
|
|
Code reference: https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py
|
|
"""
|
|
|
|
def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0,
|
|
sigma_max=1.0, sigma_min=0.003 / 1.002, max_stage=3, s_noise=1.0,
|
|
num_integration_points=200):
|
|
self.num_train_timesteps = num_train_timesteps
|
|
self.shift = shift
|
|
self.sigma_max = sigma_max
|
|
self.sigma_min = sigma_min
|
|
self.max_stage = max_stage
|
|
self.s_noise = s_noise
|
|
self.num_integration_points = num_integration_points
|
|
self.set_timesteps(num_inference_steps)
|
|
self.old_denoised = None
|
|
self.old_denoised_d = None
|
|
self.step_index = 0
|
|
|
|
def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, sigmas=None):
|
|
"""Generate the full sigma schedule (from max to min)."""
|
|
full_sigmas = torch.linspace(self.sigma_max, self.sigma_min, self.num_train_timesteps)
|
|
ss = len(full_sigmas) / num_inference_steps
|
|
if sigmas is None:
|
|
sigmas = []
|
|
for x in range(num_inference_steps):
|
|
idx = int(round(x * ss))
|
|
sigmas.append(float(full_sigmas[idx]))
|
|
sigmas.append(0.0)
|
|
self.sigmas = torch.FloatTensor(sigmas)
|
|
self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
|
|
self.timesteps = self.sigmas[:-1] * self.num_train_timesteps
|
|
self.step_index = 0
|
|
self.old_denoised = None
|
|
self.old_denoised_d = None
|
|
|
|
def default_er_sde_noise_scaler(self, x):
|
|
return x * ((x ** 0.3).exp() + 10.0)
|
|
|
|
def step(self, model_output, timestep, sample, generator):
|
|
|
|
if timestep.ndim == 2:
|
|
timestep = timestep.flatten(0, 1)
|
|
|
|
self.sigmas = self.sigmas.to(model_output.device)
|
|
self.timesteps = self.timesteps.to(model_output.device)
|
|
|
|
if timestep.ndim == 0:
|
|
timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
|
|
else:
|
|
timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
|
|
|
|
noise_scaler = self.default_er_sde_noise_scaler
|
|
|
|
# Get current and next sigma
|
|
sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
|
|
if (timestep_id + 1 >= len(self.sigmas)).any():
|
|
sigma_next = torch.zeros_like(sigma)
|
|
else:
|
|
sigma_next = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
|
|
|
|
er_lambda_s = sigma
|
|
er_lambda_t = sigma_next
|
|
|
|
# Calculate alpha values
|
|
alpha_s = sigma / (er_lambda_s + 1e-10)
|
|
alpha_t = sigma_next / (er_lambda_t + 1e-10)
|
|
r_alpha = alpha_t / (alpha_s + 1e-10)
|
|
|
|
# Denoised prediction (x_0 estimate)
|
|
denoised = sample - sigma * model_output
|
|
|
|
# Determine which stage to use
|
|
stage_used = min(self.max_stage, self.step_index + 1)
|
|
|
|
if sigma_next == 0 or (sigma_next == 0.0).all():
|
|
# Final step - return denoised
|
|
x = denoised
|
|
else:
|
|
r = noise_scaler(er_lambda_t) / (noise_scaler(er_lambda_s) + 1e-10)
|
|
|
|
# Stage 1: Euler step
|
|
x = r_alpha * r * sample + alpha_t * (1 - r) * denoised
|
|
|
|
if stage_used >= 2 and self.old_denoised is not None:
|
|
dt = er_lambda_t - er_lambda_s
|
|
lambda_step_size = -dt / self.num_integration_points
|
|
|
|
# Create integration points
|
|
point_indice = torch.arange(0, self.num_integration_points,
|
|
dtype=torch.float32, device=sample.device)
|
|
lambda_pos = er_lambda_t + point_indice * lambda_step_size
|
|
scaled_pos = noise_scaler(lambda_pos)
|
|
|
|
# Stage 2: Second-order correction
|
|
s = torch.sum(1 / (scaled_pos + 1e-10)) * lambda_step_size
|
|
|
|
# Get previous sigma for derivative calculation
|
|
if timestep_id > 0:
|
|
sigma_prev = self.sigmas[timestep_id - 1].reshape(-1, 1, 1, 1)
|
|
er_lambda_prev = sigma_prev
|
|
else:
|
|
er_lambda_prev = er_lambda_s
|
|
|
|
denoised_d = (denoised - self.old_denoised) / ((er_lambda_s - er_lambda_prev) + 1e-10)
|
|
x = x + alpha_t * (dt + s * noise_scaler(er_lambda_t)) * denoised_d
|
|
|
|
if stage_used >= 3 and self.old_denoised_d is not None:
|
|
# Stage 3: Third-order correction
|
|
s_u = torch.sum((lambda_pos - er_lambda_s) / (scaled_pos + 1e-10)) * lambda_step_size
|
|
|
|
# Get sigma from two steps ago
|
|
if timestep_id > 1:
|
|
sigma_prev_prev = self.sigmas[timestep_id - 2].reshape(-1, 1, 1, 1)
|
|
er_lambda_prev_prev = sigma_prev_prev
|
|
else:
|
|
er_lambda_prev_prev = er_lambda_prev
|
|
|
|
denoised_u = (denoised_d - self.old_denoised_d) / (((er_lambda_s - er_lambda_prev_prev) / 2) + 1e-10)
|
|
x = x + alpha_t * ((dt ** 2) / 2 + s_u * noise_scaler(er_lambda_t)) * denoised_u
|
|
|
|
self.old_denoised_d = denoised_d
|
|
|
|
# Add stochastic noise
|
|
if self.s_noise > 0:
|
|
noise_term = (er_lambda_t ** 2 - er_lambda_s ** 2 * r ** 2).sqrt()
|
|
noise_term = torch.nan_to_num(noise_term, nan=0.0)
|
|
noise = torch.randn(*x.shape, dtype=torch.float32, device=torch.device("cpu"), generator=generator).to(x)
|
|
x = x + alpha_t * noise * self.s_noise * noise_term
|
|
|
|
# Store current denoised for next iteration
|
|
self.old_denoised = denoised
|
|
self.step_index += 1
|
|
|
|
return x
|
|
|
|
def add_noise(self, original_samples, noise, timestep):
|
|
if timestep.ndim == 2:
|
|
timestep = timestep.flatten(0, 1)
|
|
|
|
self.sigmas = self.sigmas.to(noise.device)
|
|
self.timesteps = self.timesteps.to(noise.device)
|
|
|
|
timestep_id = torch.argmin(
|
|
(self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
|
|
sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
|
|
|
|
sample = (1 - sigma) * original_samples + sigma * noise
|
|
return sample.type_as(noise)
|