"""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