File size: 2,922 Bytes
101858b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import torch
import math
class LatentDiffusionScheduler:
def __init__(self, num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"):
self.num_train_timesteps = num_train_timesteps
if beta_schedule == "scaled_linear":
# standard SD schedule
self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
else:
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
timesteps = timesteps.to(original_samples.device)
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
sqrt_alpha_prod = sqrt_alpha_prod.view(-1, 1, 1, 1)
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.view(-1, 1, 1, 1)
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
return noisy_samples
def step(
self,
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
prev_timestep: int,
clip_sample: bool = True,
clip_range: float = 2.0,
) -> torch.Tensor:
device = sample.device
dtype = sample.dtype
alphas_cumprod = self.alphas_cumprod.to(device=device, dtype=dtype)
alpha_prod_t = alphas_cumprod[timestep]
alpha_prod_t_prev = alphas_cumprod[prev_timestep] if prev_timestep >= 0 else torch.tensor(1.0, device=device, dtype=dtype)
beta_prod_t = 1 - alpha_prod_t
beta_prod_t_prev = 1 - alpha_prod_t_prev
# DDIM step for epsilon prediction
# 1. Predict clean latent (x0)
pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
# 2. Optional clamp. Training supervises wider SDXL VAE latents, so
# generation can disable this to avoid washing out valid latent detail.
if clip_sample:
pred_original_sample = pred_original_sample.clamp(-clip_range, clip_range)
# 3. Compute direction pointing to xt
pred_prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + beta_prod_t_prev ** 0.5 * model_output
return pred_prev_sample
def set_timesteps(self, num_inference_steps: int, device: torch.device):
# Linear spacing of timesteps
timesteps = torch.linspace(self.num_train_timesteps - 1, 0, num_inference_steps).round().long()
return timesteps.to(device)
|