Spaces:
Running on Zero
Running on Zero
| """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() | |
| 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) | |
| 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) | |