| """Pixel-space U-Net velocity model with self+cross attention to char tokens.""" |
| import math |
| import torch, torch.nn as nn, torch.nn.functional as F |
| from conditioning import TextConditioner |
|
|
|
|
| def timestep_embedding(t, dim): |
| |
| half = dim // 2 |
| w = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device).float() / half) |
| a = (t.float() * 1000.0)[:, None] * w[None] |
| return torch.cat([torch.cos(a), torch.sin(a)], -1) |
|
|
|
|
| def gn(c): return nn.GroupNorm(32, c) |
|
|
|
|
| class ResBlock(nn.Module): |
| def __init__(self, cin, cout, tdim, dropout=0.0): |
| super().__init__() |
| self.n1, self.c1 = gn(cin), nn.Conv2d(cin, cout, 3, padding=1) |
| self.t = nn.Linear(tdim, cout * 2) |
| self.n2, self.c2 = gn(cout), nn.Conv2d(cout, cout, 3, padding=1) |
| self.drop = nn.Dropout(dropout) |
| self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity() |
| nn.init.zeros_(self.c2.weight); nn.init.zeros_(self.c2.bias) |
|
|
| def forward(self, x, temb): |
| h = self.c1(F.silu(self.n1(x))) |
| s, b = self.t(F.silu(temb))[:, :, None, None].chunk(2, 1) |
| h = self.n2(h) * (1 + s) + b |
| h = self.c2(self.drop(F.silu(h))) |
| return self.skip(x) + h |
|
|
|
|
| class AttnBlock(nn.Module): |
| """Self-attention over pixels, then cross-attention pixels->char tokens, then MLP.""" |
| def __init__(self, c, ctx_dim, heads=8): |
| super().__init__() |
| self.h = heads |
| self.n1, self.qkv, self.o1 = gn(c), nn.Linear(c, 3 * c), nn.Linear(c, c) |
| self.n2, self.q, self.kv, self.o2 = nn.LayerNorm(c), nn.Linear(c, c), nn.Linear(ctx_dim, 2 * c), nn.Linear(c, c) |
| self.n3, self.mlp = nn.LayerNorm(c), nn.Sequential(nn.Linear(c, 4 * c), nn.GELU(), nn.Linear(4 * c, c)) |
| for o in (self.o1, self.o2, self.mlp[2]): nn.init.zeros_(o.weight); nn.init.zeros_(o.bias) |
|
|
| def split(self, x): |
| B, N, C = x.shape; return x.view(B, N, self.h, C // self.h).transpose(1, 2) |
|
|
| def forward(self, x, ctx, mask): |
| B, C, H, W = x.shape |
| h = self.n1(x).flatten(2).transpose(1, 2) |
| q, k, v = self.qkv(h).chunk(3, -1) |
| a = F.scaled_dot_product_attention(self.split(q), self.split(k), self.split(v)) |
| t = x.flatten(2).transpose(1, 2) + self.o1(a.transpose(1, 2).reshape(B, -1, C)) |
| k, v = self.kv(ctx).chunk(2, -1) |
| am = mask[:, None, None, :] |
| a = F.scaled_dot_product_attention(self.split(self.q(self.n2(t))), self.split(k), self.split(v), attn_mask=am) |
| t = t + self.o2(a.transpose(1, 2).reshape(B, -1, C)) |
| t = t + self.mlp(self.n3(t)) |
| return t.transpose(1, 2).reshape(B, C, H, W) |
|
|
|
|
| class UNet(nn.Module): |
| """Levels at 128px input: 128,64,32,16 (attn at the 32 and 16 levels, i.e. levels 2,3; |
| at 256px input the same levels sit at 64/32).""" |
| def __init__(self, ch=64, mults=(1, 2, 3, 4), attn_levels=(2, 3), nres=2, ctx_dim=256, tdim=512): |
| super().__init__() |
| self.tmlp = nn.Sequential(nn.Linear(ch, tdim), nn.SiLU(), nn.Linear(tdim, tdim)) |
| self.ch = ch |
| self.inp = nn.Conv2d(3, ch, 3, padding=1) |
| self.down, chans, c = nn.ModuleList(), [ch], ch |
| for i, m in enumerate(mults): |
| for _ in range(nres): |
| blk = nn.ModuleList([ResBlock(c, ch * m, tdim)]); c = ch * m |
| if i in attn_levels: blk.append(AttnBlock(c, ctx_dim)) |
| self.down.append(blk); chans.append(c) |
| if i < len(mults) - 1: |
| self.down.append(nn.ModuleList([nn.Conv2d(c, c, 3, stride=2, padding=1)])); chans.append(c) |
| self.mid = nn.ModuleList([ResBlock(c, c, tdim), AttnBlock(c, ctx_dim), ResBlock(c, c, tdim)]) |
| self.up = nn.ModuleList() |
| for i, m in reversed(list(enumerate(mults))): |
| for _ in range(nres + 1): |
| blk = nn.ModuleList([ResBlock(c + chans.pop(), ch * m, tdim)]); c = ch * m |
| if i in attn_levels: blk.append(AttnBlock(c, ctx_dim)) |
| self.up.append(blk) |
| if i > 0: |
| self.up.append(nn.ModuleList([nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(c, c, 3, padding=1)])) |
| self.out = nn.Sequential(gn(c), nn.SiLU(), nn.Conv2d(c, 3, 3, padding=1)) |
| nn.init.zeros_(self.out[2].weight); nn.init.zeros_(self.out[2].bias) |
|
|
| def _run(self, blk, h, temb, ctx, mask): |
| for m in blk: |
| if isinstance(m, ResBlock): h = m(h, temb) |
| elif isinstance(m, AttnBlock): h = m(h, ctx, mask) |
| else: h = m(h) |
| return h |
|
|
| def forward(self, x, t, ctx, mask): |
| temb = self.tmlp(timestep_embedding(t, self.ch)) |
| h = self.inp(x); hs = [h] |
| for blk in self.down: |
| h = self._run(blk, h, temb, ctx, mask); hs.append(h) |
| h = self._run(self.mid, h, temb, ctx, mask) |
| for blk in self.up: |
| if isinstance(blk[0], ResBlock): h = torch.cat([h, hs.pop()], 1) |
| h = self._run(blk, h, temb, ctx, mask) |
| return self.out(h) |
|
|
|
|
| class FlowModel(nn.Module): |
| """Text conditioner + U-Net, trained jointly. forward returns predicted velocity.""" |
| def __init__(self, ctx_dim=256, **kw): |
| super().__init__() |
| self.cond = TextConditioner(dim=ctx_dim) |
| self.unet = UNet(ctx_dim=ctx_dim, **kw) |
|
|
| def forward(self, x, t, ids, mask, force_uncond=False): |
| ctx, m = self.cond(ids, mask, force_uncond=force_uncond) |
| return self.unet(x, t, ctx, m) |
|
|
|
|
| PARAM_RANGE = (25e6, 40e6) |
|
|
|
|
| def build_model(verbose=True): |
| m = FlowModel() |
| n = sum(p.numel() for p in m.parameters()) |
| if verbose: |
| nu = sum(p.numel() for p in m.unet.parameters()) |
| print(f"params: total {n/1e6:.2f}M (unet {nu/1e6:.2f}M, text {(n-nu)/1e6:.2f}M)") |
| assert PARAM_RANGE[0] <= n <= PARAM_RANGE[1], f"param count {n} out of range" |
| return m |
|
|
|
|
| if __name__ == "__main__": |
| from conditioning import tokenize |
| m = build_model() |
| ids, mask = tokenize(["a red smiley face with 'HI' next to it"] * 2) |
| print(m(torch.randn(2, 3, 128, 128), torch.rand(2), ids, mask).shape) |
|
|
