Spaces:
Running on Zero
Running on Zero
| """The Thinker: condition encoder -> TRM recursion over tokens -> decoder -> ControlNet adapters into a frozen Painter. | |
| TRM recursion (Jolicoeur-Martineau, 2025): | |
| deep_recursion: H cycles, the first H-1 without grad; each cycle does | |
| L times z <- net(z, y + c) (latent update, input c supplied) | |
| once y <- net(y, z) (answer update) | |
| deep supervision: N_sup calls of deep_recursion, (y, z) detached in between. | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .unet import UNet, timestep_embedding | |
| def rms_norm(x, eps=1e-5): | |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) | |
| def rope(x, cos, sin): | |
| x1, x2 = x.chunk(2, -1) | |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1) | |
| class Block(nn.Module): | |
| """Post-norm transformer block (attention + SwiGLU), as in TRM.""" | |
| def __init__(self, dim, heads, expansion): | |
| super().__init__() | |
| self.heads = heads | |
| self.qkv, self.o = nn.Linear(dim, 3 * dim, bias=False), nn.Linear(dim, dim, bias=False) | |
| inter = 256 * math.ceil(round(expansion * dim * 2 / 3) / 256) | |
| self.gate_up, self.down = nn.Linear(dim, 2 * inter, bias=False), nn.Linear(inter, dim, bias=False) | |
| def forward(self, x, cos, sin): | |
| B, N, D = x.shape | |
| q, k, v = self.qkv(x).view(B, N, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4) | |
| q, k = rope(q, cos, sin), rope(k, cos, sin) | |
| a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, D) | |
| x = rms_norm(x + self.o(a)) | |
| g, u = self.gate_up(x).chunk(2, -1) | |
| return rms_norm(x + self.down(F.silu(g) * u)) | |
| class Net(nn.Module): | |
| """The single weight-tied network used for both the z- and y-updates: net(h, injection).""" | |
| def __init__(self, dim=512, heads=8, layers=2, expansion=4, seq_len=81): | |
| super().__init__() | |
| self.blocks = nn.ModuleList([Block(dim, heads, expansion) for _ in range(layers)]) | |
| hd = dim // heads | |
| inv = 1.0 / (10000 ** (torch.arange(0, hd, 2).float() / hd)) | |
| ang = torch.outer(torch.arange(seq_len).float(), inv) | |
| self.register_buffer("cos", ang.cos(), persistent=False) | |
| self.register_buffer("sin", ang.sin(), persistent=False) | |
| def forward(self, h, injection): | |
| h = h + injection | |
| for b in self.blocks: | |
| h = b(h, self.cos, self.sin) | |
| return h | |
| class Encoder(nn.Module): | |
| """concat(x_t, cond, mask) -> conv stack -> avg-pool to grid x grid tokens -> dim, plus timestep embedding.""" | |
| def __init__(self, dim=512, hidden=(128, 256, 256), grid=9, in_ch=3): | |
| super().__init__() | |
| layers, c = [], in_ch | |
| for i, h in enumerate(hidden): | |
| layers += [nn.Conv2d(c, h, 3, stride=1 if i == 0 else 2, padding=1), nn.SiLU()] | |
| c = h | |
| self.conv, self.grid = nn.Sequential(*layers), grid | |
| self.proj = nn.Linear(c, dim) | |
| self.t = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim)) | |
| def forward(self, xt, t, cond, mask): | |
| h = F.adaptive_avg_pool2d(self.conv(torch.cat([xt, cond, mask], 1)), self.grid) | |
| tok = self.proj(h.flatten(2).transpose(1, 2)) | |
| return tok + self.t(timestep_embedding(t, tok.shape[-1]))[:, None] | |
| class Decoder(nn.Module): | |
| """y tokens -> image-resolution feature map -> one zero-initialised 1x1 adapter per Painter skip (+ mid).""" | |
| def __init__(self, painter: UNet, dim=512, grid=9, img=144, ch=32): | |
| super().__init__() | |
| self.grid, self.img = grid, img | |
| self.proj = nn.Conv2d(dim, 2 * ch, 1) | |
| self.up = nn.Sequential(nn.ConvTranspose2d(2 * ch, 2 * ch, 4, stride=4), nn.SiLU(), | |
| nn.ConvTranspose2d(2 * ch, ch, 2, stride=2), nn.SiLU(), | |
| nn.ConvTranspose2d(ch, ch, 2, stride=2), nn.SiLU(), | |
| nn.Conv2d(ch, ch, 3, padding=1), nn.SiLU()) | |
| def zero(c): | |
| conv = nn.Conv2d(ch, c, 1) | |
| nn.init.zeros_(conv.weight), nn.init.zeros_(conv.bias) | |
| return conv | |
| self.specs = painter.skip_specs | |
| self.adapters = nn.ModuleList([zero(c) for c, _ in self.specs]) | |
| self.mid_factor = painter.mid_spec[1] | |
| self.mid_adapter = zero(painter.mid_spec[0]) | |
| def forward(self, y): | |
| B, N, D = y.shape | |
| h = self.proj(y.transpose(1, 2).reshape(B, D, self.grid, self.grid)) | |
| h = self.up(h) # 16x upsample | |
| if h.shape[-1] != self.img: # token grids that do not divide the image (paper App. E) | |
| h = F.interpolate(h, size=(self.img, self.img), mode="bilinear", align_corners=False) | |
| pool = lambda f: F.avg_pool2d(h, f) if f > 1 else h | |
| skips = [a(pool(f)) for a, (_, f) in zip(self.adapters, self.specs)] | |
| return skips, self.mid_adapter(pool(self.mid_factor)) | |
| class Thinker(nn.Module): | |
| """Encoder + TRM net + decoder. The Painter is passed in at call time and never owned (stays frozen).""" | |
| def __init__(self, painter, dim=512, heads=8, layers=2, expansion=4, grid=9, | |
| enc_hidden=(128, 256, 256), H=3, L=6, img=144): | |
| super().__init__() | |
| self.H, self.L = H, L | |
| self.encoder = Encoder(dim, enc_hidden, grid) | |
| self.net = Net(dim, heads, layers, expansion, grid * grid) | |
| self.decoder = Decoder(painter, dim, grid, img) | |
| self.register_buffer("y_init", nn.init.trunc_normal_(torch.empty(dim), std=1.0)) | |
| self.register_buffer("z_init", nn.init.trunc_normal_(torch.empty(dim), std=1.0)) | |
| def init_carry(self, c): | |
| return self.y_init.expand_as(c).clone(), self.z_init.expand_as(c).clone() | |
| def deep_recursion(self, c, y, z): | |
| with torch.no_grad(): | |
| for _ in range(self.H - 1): | |
| for _ in range(self.L): | |
| z = self.net(z, y + c) | |
| y = self.net(y, z) | |
| for _ in range(self.L): | |
| z = self.net(z, y + c) | |
| y = self.net(y, z) | |
| return (y.detach(), z.detach()), y | |
| def residuals(self, y): | |
| return self.decoder(y) | |
| def think(thinker, xt, t, cond, mask, carry=None, n_sup=1): | |
| """Run n_sup supervision steps (no grad except the last). Returns (Painter residuals, carry).""" | |
| c = thinker.encoder(xt, t, cond, mask) | |
| y, z = carry if carry is not None else thinker.init_carry(c) | |
| with torch.no_grad(): | |
| for _ in range(n_sup - 1): | |
| (y, z), _ = thinker.deep_recursion(c, y, z) | |
| carry, y_out = thinker.deep_recursion(c, y, z) | |
| return thinker.residuals(y_out), carry | |
| def path_forward(thinker, painter, xt, t, cond, mask, carry=None, n_sup=1): | |
| """Thinker then Painter. Returns (x0_hat, carry).""" | |
| residuals, carry = think(thinker, xt, t, cond, mask, carry, n_sup) | |
| return painter(xt, t, residuals), carry | |