lars
8cebc4809d
First repo-wide ruff format pass, plus a note in CLAUDE.md to run ruff and ty periodically. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
import torch
|
|
|
|
from giant.constants import X_DIM
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_flow(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
steps: int = 10,
|
|
) -> torch.Tensor:
|
|
"""Euler integration of the learned vector field from t=0 to t=1."""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
dt = 1.0 / steps
|
|
for i in range(steps):
|
|
t = torch.full((B,), i * dt, device=device)
|
|
v = model(x, t, cond_cont, cond_cat)
|
|
x = x + v * dt
|
|
return x
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_ddpm(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
schedule,
|
|
) -> torch.Tensor:
|
|
"""Full DDPM ancestral sampling (T reverse steps)."""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
T = schedule.T
|
|
for i in reversed(range(T)):
|
|
t_norm = torch.full((B,), i / T, device=device)
|
|
eps_pred = model(x, t_norm, cond_cont, cond_cat)
|
|
beta = schedule.betas[i]
|
|
alpha = schedule.alphas[i]
|
|
alpha_bar = schedule.alpha_bars[i]
|
|
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
|
|
x = (1.0 / alpha.sqrt()) * (
|
|
x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
|
|
) + beta.sqrt() * z
|
|
return x
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_ddim(
|
|
model: torch.nn.Module,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
schedule,
|
|
steps: int = 50,
|
|
) -> torch.Tensor:
|
|
"""DDIM deterministic sampling (Song et al. 2020) with `steps` substeps."""
|
|
model.eval()
|
|
B = cond_cont.size(0)
|
|
device = cond_cont.device
|
|
T = schedule.T
|
|
timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
|
|
x = torch.randn(B, X_DIM, device=device)
|
|
for step_idx, ts in enumerate(timesteps):
|
|
t_idx = int(ts.item())
|
|
t_norm = torch.full((B,), t_idx / T, device=device)
|
|
eps_pred = model(x, t_norm, cond_cont, cond_cat)
|
|
ab_t = schedule.alpha_bars[t_idx]
|
|
if step_idx + 1 < len(timesteps):
|
|
ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
|
|
else:
|
|
ab_prev = torch.ones(1, device=device)
|
|
x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
|
|
x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
|
|
return x
|