| """HoLo-FuSe inference: UNet + frozen-HSL label conditioning + respaced DDIM sampler. |
| Matches the training code in https://github.com/Woojiggun/HoLo-FuSe (experiment.py). |
| """ |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import hsl_embedding_zero as hslz |
|
|
| FEAT_DIM = hslz.FEAT_DIM |
|
|
|
|
| def timestep_emb(t, dim): |
| half = dim // 2 |
| freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / (half - 1)) |
| a = t.float()[:, None] * freqs[None] |
| return torch.cat([torch.sin(a), torch.cos(a)], dim=-1) |
|
|
|
|
| class ResBlock(nn.Module): |
| def __init__(self, cin, cout, emb_dim, groups=8): |
| super().__init__() |
| self.n1 = nn.GroupNorm(groups, cin); self.c1 = nn.Conv2d(cin, cout, 3, padding=1) |
| self.emb = nn.Linear(emb_dim, cout) |
| self.n2 = nn.GroupNorm(groups, cout); self.c2 = nn.Conv2d(cout, cout, 3, padding=1) |
| self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity() |
|
|
| def forward(self, x, emb): |
| h = self.c1(F.silu(self.n1(x))) |
| h = h + self.emb(emb)[:, :, None, None] |
| h = self.c2(F.silu(self.n2(h))) |
| return h + self.skip(x) |
|
|
|
|
| class AttnBlock(nn.Module): |
| def __init__(self, ch, groups=8): |
| super().__init__() |
| self.norm = nn.GroupNorm(groups, ch) |
| self.qkv = nn.Conv2d(ch, ch * 3, 1) |
| self.proj = nn.Conv2d(ch, ch, 1) |
|
|
| def forward(self, x): |
| B, C, H, W = x.shape |
| q, k, v = self.qkv(self.norm(x)).chunk(3, dim=1) |
| q = q.reshape(B, C, H * W).permute(0, 2, 1) |
| k = k.reshape(B, C, H * W) |
| v = v.reshape(B, C, H * W).permute(0, 2, 1) |
| a = torch.softmax(q @ k * (C ** -0.5), dim=-1) |
| o = (a @ v).permute(0, 2, 1).reshape(B, C, H, W) |
| return x + self.proj(o) |
|
|
|
|
| class Downsample(nn.Module): |
| def __init__(self, ch): |
| super().__init__(); self.op = nn.Conv2d(ch, ch, 3, stride=2, padding=1) |
|
|
| def forward(self, x, emb=None): return self.op(x) |
|
|
|
|
| class Upsample(nn.Module): |
| def __init__(self, ch): |
| super().__init__(); self.op = nn.Conv2d(ch, ch, 3, padding=1) |
|
|
| def forward(self, x, emb=None): return self.op(F.interpolate(x, scale_factor=2, mode="nearest")) |
|
|
|
|
| class UNet(nn.Module): |
| def __init__(self, c=3, base=128, ch_mults=(1, 2, 2, 2), num_res=2, attn_res=(16,), |
| emb_dim=256, cond_dim=128, img_size=128): |
| super().__init__() |
| self.emb_dim = emb_dim |
| attn_res = set(attn_res) |
| self.tproj = nn.Sequential(nn.Linear(emb_dim, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim)) |
| self.cproj = nn.Linear(cond_dim, emb_dim) |
| self.in_conv = nn.Conv2d(c, base, 3, padding=1) |
| L = len(ch_mults) |
| chs = [base]; now = base |
| self.downs = nn.ModuleList() |
| for i, m in enumerate(ch_mults): |
| res = img_size >> i; out = base * m |
| for _ in range(num_res): |
| grp = [ResBlock(now, out, emb_dim)]; now = out |
| if res in attn_res: grp.append(AttnBlock(out)) |
| self.downs.append(nn.ModuleList(grp)); chs.append(now) |
| if i != L - 1: |
| self.downs.append(Downsample(now)); chs.append(now) |
| self.mid1 = ResBlock(now, now, emb_dim); self.mid_attn = AttnBlock(now); self.mid2 = ResBlock(now, now, emb_dim) |
| self.ups = nn.ModuleList() |
| for i, m in reversed(list(enumerate(ch_mults))): |
| res = img_size >> i; out = base * m |
| for _ in range(num_res + 1): |
| grp = [ResBlock(now + chs.pop(), out, emb_dim)]; now = out |
| if res in attn_res: grp.append(AttnBlock(out)) |
| self.ups.append(nn.ModuleList(grp)) |
| if i != 0: |
| self.ups.append(Upsample(now)) |
| self.out = nn.Sequential(nn.GroupNorm(8, now), nn.SiLU(), nn.Conv2d(now, c, 3, padding=1)) |
|
|
| def forward(self, x, t, cond): |
| emb = self.tproj(timestep_emb(t, self.emb_dim)) |
| if cond is not None: |
| emb = emb + self.cproj(cond) |
| h = self.in_conv(x); hs = [h] |
| for stage in self.downs: |
| if isinstance(stage, Downsample): |
| h = stage(h) |
| else: |
| for blk in stage: |
| h = blk(h, emb) if isinstance(blk, ResBlock) else blk(h) |
| hs.append(h) |
| h = self.mid1(h, emb); h = self.mid_attn(h); h = self.mid2(h, emb) |
| for stage in self.ups: |
| if isinstance(stage, Upsample): |
| h = stage(h) |
| else: |
| h = torch.cat([h, hs.pop()], dim=1) |
| for blk in stage: |
| h = blk(h, emb) if isinstance(blk, ResBlock) else blk(h) |
| return self.out(h) |
|
|
|
|
| class HSLLabelCond(nn.Module): |
| """label bytes -> frozen 27-D HSL frame (0 learned params) -> small learned readout.""" |
| def __init__(self, cond_dim, K=8): |
| super().__init__() |
| self.K = K |
| self.zero = hslz.ZeroInput(K=K, dim=512) |
| self.readout = nn.Sequential(nn.Linear(K * FEAT_DIM, cond_dim), nn.SiLU(), |
| nn.Linear(cond_dim, cond_dim)) |
|
|
| @torch.no_grad() |
| def _hsl(self, label_strs): |
| ids = [] |
| for s in label_strs: |
| b = s.encode("utf-8", "replace")[:self.K] |
| b = b + b"\x00" * (self.K - len(b)) |
| ids.append(list(b)) |
| ids = torch.tensor(ids, dtype=torch.long) |
| return self.zero.features(ids).reshape(len(label_strs), -1) |
|
|
| def forward(self, label_strs): |
| return self.readout(self._hsl(label_strs).float().to(next(self.parameters()).device)) |
|
|
|
|
| class LearnedLabelCond(nn.Module): |
| """Control arm: same-budget learned embedding (definition matches training exactly).""" |
| def __init__(self, cond_dim, class_list): |
| super().__init__() |
| self.idx = {c: i for i, c in enumerate(class_list)} |
| self.emb = nn.Embedding(len(class_list), cond_dim) |
| self.readout = nn.Sequential(nn.Linear(cond_dim, cond_dim), nn.SiLU(), |
| nn.Linear(cond_dim, cond_dim)) |
|
|
| def forward(self, label_strs): |
| ii = torch.tensor([self.idx.get(s, 0) for s in label_strs], |
| device=self.emb.weight.device) |
| return self.readout(self.emb(ii)) |
|
|
|
|
| def make_cosine_acp(T=250): |
| s = 0.008 |
| steps = torch.arange(T + 1, dtype=torch.float64) |
| ac = torch.cos(((steps / T) + s) / (1 + s) * math.pi / 2) ** 2 |
| ac = ac / ac[0] |
| betas = (1 - ac[1:] / ac[:-1]).clamp(1e-4, 0.999).float() |
| return torch.cumprod(1 - betas, 0) |
|
|
|
|
| @torch.no_grad() |
| def ddim_sample(model, cond, uncond, n=1, c=3, hw=128, T=250, steps=16, cfg=1.6, |
| dyn=0.99, seed=None, progress=None): |
| """Respaced DDIM (eta=0) over the trained T=250 cosine schedule, with CFG + dynamic thresholding.""" |
| dev = next(model.parameters()).device |
| if seed is not None and seed >= 0: |
| torch.manual_seed(seed) |
| acp = make_cosine_acp(T) |
| ts = torch.linspace(0, T - 1, steps).round().long().unique().tolist() |
| x = torch.randn(n, c, hw, hw, device=dev) |
| for k in range(len(ts) - 1, -1, -1): |
| t = ts[k] |
| tt = torch.full((n,), t, dtype=torch.long, device=dev) |
| if cond is not None and cfg != 1.0: |
| eps = model(x, tt, uncond) if uncond is not None else model(x, tt, None) |
| eps = eps + cfg * (model(x, tt, cond) - eps) |
| else: |
| eps = model(x, tt, cond) |
| a = acp[t] |
| x0 = (x - (1 - a).sqrt() * eps) / a.sqrt() |
| if dyn and dyn > 0: |
| s = torch.quantile(x0.reshape(n, -1).abs(), dyn, dim=1).clamp(min=1.0).view(-1, 1, 1, 1) |
| x0 = x0.clamp(-s, s) / s |
| else: |
| x0 = x0.clamp(-1, 1) |
| if k == 0: |
| x = x0 |
| else: |
| ap = acp[ts[k - 1]] |
| x = ap.sqrt() * x0 + (1 - ap).sqrt() * eps |
| if progress is not None: |
| progress((len(ts) - k) / len(ts)) |
| return x.clamp(-1, 1) |
|
|
|
|
| def load_holofuse(arm="hsl", ckpt_path=None, device="cpu"): |
| """One-call loader: returns (model, cond_enc) with EMA weights applied. |
| arm: 'hsl' | 'learned' | 'none'. Downloads from ggunio/HoLo-FuSe unless ckpt_path is given.""" |
| if ckpt_path is None: |
| from huggingface_hub import hf_hub_download |
| ckpt_path = hf_hub_download("ggunio/HoLo-FuSe", f"holofuse_{arm}_128.pt") |
| st = torch.load(ckpt_path, map_location=device, weights_only=False) |
| model = UNet().to(device).eval() |
| if arm == "hsl": |
| cond_enc = HSLLabelCond(128).to(device).eval() |
| elif arm == "learned": |
| cond_enc = LearnedLabelCond(128, ["Cat", "Dog"]).to(device).eval() |
| else: |
| cond_enc = None |
| params = list(model.parameters()) + (list(cond_enc.parameters()) if cond_enc else []) |
| with torch.no_grad(): |
| for p, e in zip(params, st["ema"]): |
| p.copy_(e.to(device)) |
| return model, cond_enc |
|
|
|
|
| def generate(label="Cat", arm="hsl", steps=16, cfg=1.6, seed=0, n=1, device="cpu", |
| model=None, cond_enc=None): |
| """One-call generation -> list of PIL images (128px).""" |
| from PIL import Image |
| if model is None: |
| model, cond_enc = load_holofuse(arm, device=device) |
| with torch.no_grad(): |
| cond = cond_enc([label] * n) if cond_enc else None |
| uncond = torch.zeros_like(cond) if cond is not None else None |
| x = ddim_sample(model, cond, uncond, n=n, steps=steps, cfg=cfg, dyn=0.99, seed=seed) |
| arrs = ((x.clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu().numpy().transpose(0, 2, 3, 1) |
| return [Image.fromarray(a, "RGB") for a in arrs] |
|
|
