Spaces:
Running on Zero
Running on Zero
File size: 2,547 Bytes
65635a9 | 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 | """x0-parameterised diffusion: cosine schedule, forward noising, DDIM sampling with CFG, EMA."""
import copy
import math
import torch
class Schedule:
def __init__(self, num_train_timesteps=100):
self.T = num_train_timesteps
s = 0.008
f = lambda t: math.cos((t / self.T + s) / (1 + s) * math.pi / 2) ** 2
betas = [min(1 - f(t + 1) / f(t), 0.999) for t in range(self.T)]
self.alpha_bar = torch.cumprod(1 - torch.tensor(betas, dtype=torch.float64), 0).float()
def ab(self, t, device):
return self.alpha_bar.to(device)[t].view(-1, 1, 1, 1)
def q_sample(self, x0, t, noise):
ab = self.ab(t, x0.device)
return ab.sqrt() * x0 + (1 - ab).sqrt() * noise
def sample_t(self, n, device):
return torch.randint(0, self.T, (n,), device=device)
def timesteps(self, steps):
"""Evenly spaced descending timesteps, e.g. 99, 94, ..., 0 for steps=20."""
return torch.linspace(self.T - 1, 0, steps).round().long().tolist()
@torch.no_grad()
def ddim(self, model_fn, x, steps=20, start=None, callback=None):
"""Deterministic DDIM. model_fn(x_t, t_tensor, i) -> x0_hat. Returns x_0.
`start` lets sampling begin from an intermediate step index (used for re-noising diagnostics).
`callback(i, t, x0_hat, x_t)` is called at every step with the prediction and the current noisy sample.
"""
ts = self.timesteps(steps)
for i, t in enumerate(ts):
if start is not None and i < start:
continue
tt = torch.full((x.shape[0],), t, device=x.device, dtype=torch.long)
x0 = model_fn(x, tt, i).clamp(-1, 1)
if callback is not None:
callback(i, t, x0, x)
ab = self.ab(tt, x.device)
ab_prev = self.ab(torch.full_like(tt, ts[i + 1]), x.device) if i + 1 < len(ts) else torch.ones_like(ab)
eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
x = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps
return x
def cfg(cond_out, uncond_out, scale):
return uncond_out + scale * (cond_out - uncond_out)
class EMA:
def __init__(self, model, rate=0.999):
self.rate = rate
self.model = copy.deepcopy(model).eval().requires_grad_(False)
@torch.no_grad()
def update(self, model):
torch._foreach_lerp_(list(self.model.parameters()), list(model.parameters()), 1 - self.rate)
for e, b in zip(self.model.buffers(), model.buffers()):
e.copy_(b)
|