"""The shipped Idol graph set: GPU-shaped modules built from the real weights, one module per LiteRT graph. Every GPU call stays well under 2 s on the S26 (CONTRACT.md "Forms"), so the big parts are split: denoiser den_ctx (condition refiner, once per LoRA phase), den_in (time embedding, noise refiner, camera MLP), den_joint_00 .. den_joint_23 (one joint block each; _00 concatenates the streams, _23 adds the heads) octree oct (one graph, one signature per row bucket 64 / 256 / 1024 / 4096 / 8192: the weights are shared) gaussians gs_0 .. gs_7 (two blocks each; gs_0 embeds the anchors, gs_7 projects to the 480 features) dinov3 dino_0 .. dino_3 (eight layers each; dino_0 normalises and embeds the pixels, dino_3 ends in both norms) vae vae_0 .. vae_2 (one down stage each), vae_3 (down 3 + mid resnet 0), vae_attn (mid attention, fp32), vae_4 (mid resnet 1, head, quant_conv) Host inputs instead of in-graph math, for fp16 range: sinusoidal features of the diffusion time and the octree level (args up to 1000 rad), the point position embeddings (args up to 2^16 * 2pi), the denoiser's fixed latent-token position embedding, and DINOv3's RoPE tables. The turbo LoRA is kept as side matmuls: one graph set, and a gate input [1, 64] = (4 * slow, 4 * fast) per 32 ranks selects the phase. """ import math import torch import torch.nn as nn import torch.nn.functional as F import gpu_blocks as G LORA_SCALE = 128.0 / 32.0 DEN_TOKENS, COND_TOKENS = 8192, 4101 OCT_BUCKETS = (64, 256, 1024, 4096, 8192) def lora_pair(sd, key): """-> A [64, in] (slow rows then fast rows), B [out, 64] (slow columns then fast columns).""" a = torch.cat([sd[f'{key}.slow_A'], sd[f'{key}.fast_A']], 0).float() b = torch.cat([sd[f'{key}.slow_B'], sd[f'{key}.fast_B']], 1).float() return a, b def lora_gate(phase): g = torch.zeros(1, 64) g[:, :32] = LORA_SCALE if phase == 'slow' else 0.0 g[:, 32:] = LORA_SCALE if phase == 'fast' else 0.0 return g class DenBlock(G.DenoiserBlockGPU): """DenoiserBlockGPU + the turbo LoRA on qkv, out, fc1, fc2: forward(x, mod or None, gate [1, 64]).""" def __init__(self, block, repo, lora_sd, key, chunk): super().__init__(block, repo, chunk=chunk) c, h, d = self.c, self.h, self.d perm = G.deinterleave_perm(d) a, b = lora_pair(lora_sd, f'{key}.attn.qkv') self.qkv_a = G._linear(a, None) def rows(i, permute): bi = b[i * c:(i + 1) * c].reshape(h, d, 64) if permute: bi = bi[:, perm] return G._linear(bi.reshape(c, 64), None) self.q_b, self.k_b, self.v_b = rows(0, True), rows(1, True), rows(2, False) self.out_a, self.out_b = [G._linear(t, None) for t in lora_pair(lora_sd, f'{key}.attn.out')] self.fc1_a, self.fc1_b = [G._linear(t, None) for t in lora_pair(lora_sd, f'{key}.mlp.mlp.0')] self.fc2_a, self.fc2_b = [G._linear(t, None) for t in lora_pair(lora_sd, f'{key}.mlp.mlp.2')] def attn_l(self, hh, cos, sin, gate): b, l, _ = hh.shape half = self.d // 2 def rope(t): t0, t1 = t[..., :half], t[..., half:] return torch.cat([t0 * cos - t1 * sin, t0 * sin + t1 * cos], dim=-1) ax = self.qkv_a(hh) * gate q = G.safe_rms_unit(rope((self.q(hh) + self.q_b(ax)).reshape(b, l, self.h, self.d))) * self.gq k = G.safe_rms_unit(rope((self.k(hh) + self.k_b(ax)).reshape(b, l, self.h, self.d))) * self.gk v = (self.v(hh) + self.v_b(ax)).reshape(b, l, self.h, self.d) o = G.attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.chunk) o = o.reshape(b, l, self.c) return self.out(o) + self.out_b(self.out_a(o) * gate) def mlp_l(self, hh, gate): u = self.fc1(hh) + self.fc1_b(self.fc1_a(hh) * gate) u = G._gelu_tanh(u) return self.fc2(u) + self.fc2_b(self.fc2_a(u) * gate) def forward(self, x, mod, gate): cos, sin = self.repo(x) if not self.modulation: x = x + self.attn_l(G.safe_layer_norm(x, self.eps1, self.n1w, self.n1b), cos, sin, gate) return x + self.mlp_l(G.safe_layer_norm(x, self.eps2, self.n2w, self.n2b), gate) m = mod + self.shift_table c = self.c shift_msa, scale_msa, gate_msa = m[:, 0:c], m[:, c:2 * c], m[:, 2 * c:3 * c] shift_mlp, scale_mlp, gate_mlp = m[:, 3 * c:4 * c], m[:, 4 * c:5 * c], m[:, 5 * c:6 * c] hh = G.safe_layer_norm(x, self.eps1) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) x = x + self.attn_l(hh, cos, sin, gate) * gate_msa.unsqueeze(1) hh = G.safe_layer_norm(x, self.eps2) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) return x + self.mlp_l(hh, gate) * gate_mlp.unsqueeze(1) def _mlp_seq(seq): """nn.Sequential(Linear, SiLU or GELU(tanh), Linear) -> a float copy.""" return nn.Sequential(*[G._linear(m.weight, m.bias) if isinstance(m, nn.Linear) else m for m in seq]) class DenCtx(nn.Module): """(feature1 [1, 4101, 1280], feature2 [1, 4101, 128], gate) -> h_cond [1, 4101, 1024].""" def __init__(self, den, lora_sd, chunk): super().__init__() self.e1 = G._linear(den.cond_embedder.weight, den.cond_embedder.bias) self.e2 = G._linear(den.cond_embedder2.weight, den.cond_embedder2.bias) self.blocks = nn.ModuleList([DenBlock(b.float(), r.float(), lora_sd, f'context_refiner.{i}', chunk) for i, (b, r) in enumerate(zip(den.context_refiner, den.context_repo_layers))]) def forward(self, feature1, feature2, gate): h = self.e1(feature1) + self.e2(feature2) for blk in self.blocks: h = blk(h, None, gate) return h class DenIn(nn.Module): """(latent [1, 8192, 16], camera [1, 1, 5], t_freq [1, 256], pos_emb [1, 8192, 1024], gate) -> h_x [1, 8192, 1024], h_cam [1, 1, 1024], t_mod [1, 6144], t_emb [1, 1024].""" def __init__(self, den, lora_sd, chunk): super().__init__() self.t_mlp = _mlp_seq(den.t_embedder.mlp) self.ada = _mlp_seq(den.adaLN_modulation) self.inp = G._linear(den.input_layer.weight, den.input_layer.bias) self.blocks = nn.ModuleList([DenBlock(b.float(), r.float(), lora_sd, f'noise_refiner.{i}', chunk) for i, (b, r) in enumerate(zip(den.noise_refiner, den.noise_repo_layers))]) self.cam = _mlp_seq(den.cam_refiner.mlp) def forward(self, latent, camera, t_freq, pos_emb, gate): t_emb = self.t_mlp(t_freq) t_mod = self.ada(t_emb) h = self.inp(latent) + pos_emb for blk in self.blocks: h = blk(h, t_mod, gate) return h, self.cam(camera), t_mod, t_emb class DenJoint(nn.Module): """Joint block k. first: (h_x, h_cond, h_cam, t_mod, gate) -> h; middle: (h, t_mod, gate) -> h; last: (h, t_mod, t_emb, gate) -> v_latent [1, 8192, 16], v_camera [1, 1, 5].""" def __init__(self, den, k, lora_sd, chunk): super().__init__() self.first, self.last = k == 0, k == len(den.blocks) - 1 self.block = DenBlock(den.blocks[k].float(), den.repo_layers[k].float(), lora_sd, f'blocks.{k}', chunk) if self.last: self.shift_table = nn.Parameter(den.shift_table.detach().float().clone(), requires_grad=False) self.out = G._linear(den.out_layer.weight, den.out_layer.bias) self.cam_out = G._linear(den.cam_out_layer.weight, den.cam_out_layer.bias) def forward(self, *args): if self.first: h_x, h_cond, h_cam, t_mod, gate = args h = torch.cat([h_x, h_cond, h_cam], dim=1) elif self.last: h, t_mod, t_emb, gate = args else: h, t_mod, gate = args h = self.block(h, t_mod, gate) if not self.last: return h st = self.shift_table + t_emb.unsqueeze(1) shift, scale = st[:, 0:1], st[:, 1:2] hx = G.safe_layer_norm(h[:, :DEN_TOKENS], 1e-5) * (1 + scale) + shift hc = G.safe_layer_norm(h[:, -1:], 1e-5) * (1 + scale) + shift return self.out(hx), self.cam_out(hc) class Octree(nn.Module): """(cells [1, L, 3], cells_pe [1, L, 1024], l_freq [1, 256], latent [1, 8192, 16]) -> logits [1, L, 8].""" def __init__(self, dec, chunk): super().__init__() self.in_proj = G._linear(dec.in_proj.weight, dec.in_proj.bias) self.inp = G._linear(dec.input_layer.weight, dec.input_layer.bias) self.l_mlp = _mlp_seq(dec.l_embedder.mlp) self.ada = _mlp_seq(dec.adaLN_modulation) self.blocks = nn.ModuleList([G.OctreeBlockGPU(b.float(), chunk) for b in dec.blocks]) self.out = G._linear(dec.out_proj.weight, dec.out_proj.bias) def forward(self, cells, cells_pe, l_freq, latent): h = self.inp(self.in_proj(cells) + cells_pe) mod = self.ada(self.l_mlp(l_freq)) for blk in self.blocks: h = blk(h, mod, latent) return self.out(G.safe_layer_norm(h, 1e-5)) class Gauss(nn.Module): """Blocks [lo, hi) of the Gaussian decoder. first: (anchors [1, A, 3], anchors_pe [1, A, 1024], latent) -> h; others: (h, latent) -> h; the last ends in the output projection -> features [1, A, 480].""" def __init__(self, dgs, lo, hi, chunk): super().__init__() self.first, self.last = lo == 0, hi == len(dgs.blocks) if self.first: self.in_proj = G._linear(dgs.in_proj.weight, dgs.in_proj.bias) self.inp = G._linear(dgs.input_layer.weight, dgs.input_layer.bias) self.blocks = nn.ModuleList([G.GaussianBlockGPU(dgs.blocks[i].float(), chunk) for i in range(lo, hi)]) if self.last: self.out = G._linear(dgs.out_proj.weight, dgs.out_proj.bias) def forward(self, *args): if self.first: anchors, anchors_pe, latent = args h = self.inp(self.in_proj(anchors) + anchors_pe) else: h, latent = args for blk in self.blocks: h = blk(h, latent) return self.out(G.safe_layer_norm(h, 1e-5)) if self.last else h IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) class Dino(nn.Module): """DINOv3 layers [lo, hi). first: (pixels [1, 3, 1024, 1024] in [0, 1], cos, sin) -> h [1, 4101, 1280]; others: (h, cos, sin) -> h; the last ends in the model norm and idol's F.layer_norm -> feature1.""" def __init__(self, model, lo, hi, chunk): super().__init__() self.first, self.last = lo == 0, hi == len(model.model.layer) emb = model.embeddings if self.first: self.register_buffer('mean', torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)) self.register_buffer('inv_std', 1.0 / torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)) pe = emb.patch_embeddings self.patch = nn.Conv2d(3, pe.out_channels, pe.kernel_size, pe.stride, bias=True) self.patch.weight.data = pe.weight.detach().float().clone() self.patch.bias.data = pe.bias.detach().float().clone() self.cls = nn.Parameter(emb.cls_token.detach().float().clone(), requires_grad=False) self.reg = nn.Parameter(emb.register_tokens.detach().float().clone(), requires_grad=False) self.layers = nn.ModuleList([G.DinoBlockGPU(model.model.layer[i].float(), 20, chunk) for i in range(lo, hi)]) if self.last: self.nw = nn.Parameter(model.norm.weight.detach().float().clone(), requires_grad=False) self.nb = nn.Parameter(model.norm.bias.detach().float().clone(), requires_grad=False) self.neps = model.norm.eps def forward(self, x, cos, sin): if self.first: p = self.patch((x - self.mean) * self.inv_std) # [1, 1280, 64, 64] p = p.flatten(2).transpose(1, 2) x = torch.cat([self.cls, self.reg, p], dim=1) for layer in self.layers: x = layer(x, cos, sin) if self.last: x = G.safe_layer_norm(x, self.neps, self.nw, self.nb) x = G.safe_layer_norm(x, 1e-5) return x # -- FLUX.2 VAE encoder ---------------------------------------------------------------------------------------- class PooledGroupNorm(nn.Module): """GroupNorm without moving channels into the batch: per-channel spatial means and abs-maxes by a hierarchy of pools (no reduction over more than 1024 values at once), group statistics on [1, C] vectors, and the centred map divided by its per-group abs-max before squaring so the fp16 squares cannot overflow. Exact in real arithmetic. The abs-max goes through a max pool first (exact: a max of maxes): one REDUCE_MAX over the 1024 x 1024 map of the VAE's first stage made each call of that graph 8 s on the Galaxy S26 GPU (vae_1 at 512 x 512: 1.6 s).""" def __init__(self, gn, hw): super().__init__() self.g, self.c, self.eps = gn.num_groups, gn.num_channels, gn.eps self.w = nn.Parameter(gn.weight.detach().float().clone().view(1, -1, 1, 1), requires_grad=False) self.b = nn.Parameter(gn.bias.detach().float().clone().view(1, -1, 1, 1), requires_grad=False) self.k = max(1, hw // 32) # group -> channel expansion as a one-hot FULLY_CONNECTED (an expand or a broadcast MUL by ones becomes # BROADCAST_TO, which has no GPU op); exact: each output is 1 * one group value self.expand = nn.Linear(self.g, self.c, bias=False) self.expand.weight.data = (torch.arange(self.c)[:, None] // (self.c // self.g) == torch.arange(self.g)[None, :]).float() def _chan_mean(self, x): if self.k > 1: x = F.avg_pool2d(x, self.k) return x.mean(dim=(2, 3), keepdim=True) # [1, C, 1, 1] def _group(self, v): # [1, C, 1, 1] -> group mean, per channel g = v.reshape(1, self.g, self.c // self.g).mean(-1) return self.expand(g).reshape(1, self.c, 1, 1) def forward(self, x): d = x - self._group(self._chan_mean(x)) m = d.abs() if self.k > 1: m = F.max_pool2d(m, self.k) m = m.amax(dim=(2, 3), keepdim=True) s = m.reshape(1, self.g, self.c // self.g).amax(-1) s = self.expand(s).reshape(1, self.c, 1, 1).clamp(min=1e-2) ds = d / s var = self._group(self._chan_mean(ds * ds)) return ds * torch.rsqrt(var + self.eps / (s * s)) * self.w + self.b def _patch_gn(module, hw): for name, child in list(module.named_children()): if isinstance(child, nn.GroupNorm): setattr(module, name, PooledGroupNorm(child, hw)) else: _patch_gn(child, hw) class VaeAttn(nn.Module): """diffusers Attention of the VAE mid block (one head, residual, rescale 1) with chunked rank-4 attention.""" def __init__(self, attn, chunk, hw): super().__init__() self.gn = PooledGroupNorm(attn.group_norm, hw) self.q, self.k, self.v = [G._linear(m.weight, m.bias) for m in (attn.to_q, attn.to_k, attn.to_v)] self.o = G._linear(attn.to_out[0].weight, attn.to_out[0].bias) self.heads, self.scale, self.chunk = attn.heads, attn.scale, chunk self.residual, self.rescale = attn.residual_connection, attn.rescale_output_factor def forward(self, x, *args, **kwargs): b, c, hh, ww = x.shape h = self.gn(x).reshape(b, c, hh * ww).transpose(1, 2) d = c // self.heads q = (self.q(h) * self.scale).reshape(b, hh * ww, self.heads, d).permute(0, 2, 1, 3) k = self.k(h).reshape(b, hh * ww, self.heads, d).permute(0, 2, 1, 3) v = self.v(h).reshape(b, hh * ww, self.heads, d).permute(0, 2, 1, 3) o = G.attention(q, k, v, self.chunk).reshape(b, hh * ww, c) o = self.o(o).transpose(1, 2).reshape(b, c, hh, ww) if self.residual: o = o + x return o / self.rescale class Vae(nn.Module): """FLUX.2 VAE encoder, part p: 0: x in [0, 1] [1, 3, 1024, 1024] -> conv_in, down 0 -> [1, 128, 512, 512]; 1: down 1 -> [1, 256, 256, 256]; 2: down 2 -> [1, 512, 128, 128]; 3: down 3, mid resnet 0 -> [1, 512, 128, 128]; 'attn': the mid block's attention (one head over 16,384 tokens, its own graph to keep each GPU call short); 4: mid resnet 1, norm, SiLU, conv_out, quant_conv -> moments [1, 64, 128, 128].""" SIZES = {0: 1024, 1: 512, 2: 256, 3: 128, 'attn': 128, 4: 128} PARTS = (0, 1, 2, 3, 'attn', 4) def __init__(self, vae, part, chunk): super().__init__() enc = vae.encoder self.part = part hw = self.SIZES[part] if part in (0, 1, 2, 3): self.down = enc.down_blocks[part] _patch_gn(self.down, hw) if part == 0: self.conv_in = enc.conv_in mid = enc.mid_block if part == 3: self.res = mid.resnets[0] _patch_gn(self.res, hw) if part == 'attn': self.attn = VaeAttn(mid.attentions[0], chunk, hw) if part == 4: self.res = mid.resnets[1] _patch_gn(self.res, hw) self.norm_out = PooledGroupNorm(enc.conv_norm_out, hw) self.conv_out = enc.conv_out self.quant_conv = vae.quant_conv def forward(self, x): p = self.part if p == 0: x = self.conv_in(x * 2 - 1) if p in (0, 1, 2, 3): x = self.down(x) if p == 3: x = self.res(x, None) if p == 'attn': x = self.attn(x) if p == 4: x = self.res(x, None) x = self.quant_conv(self.conv_out(F.silu(self.norm_out(x)))) return x # -- host-side features (float32 on the CPU; the Kotlin glue computes the same) ---------------------------------- def sinusoid(value, dim=256, max_period=10000.0, two_pi=False): """SinusoidalEmbedder's frequency features of one scalar -> [1, dim] float32.""" half = dim // 2 freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) args = torch.tensor([float(value)], dtype=torch.float32)[:, None] * freqs[None] if two_pi: args = args * 2 * torch.pi return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) def point_pe(points, channels=1024, max_res=10): """PcdAbsolutePositionEmbedderv2 of points [N, 3] in [0, 1] -> [N, 1024] float32.""" freq_dim = channels // 3 // 2 freqs = torch.pow(2.0, torch.linspace(0.0, float(max_res), steps=freq_dim, dtype=torch.float32)) ang = points.float().unsqueeze(-1) * freqs * torch.pi emb = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1).reshape(points.shape[0], -1) return torch.cat([emb, torch.zeros(points.shape[0], channels - emb.shape[1])], dim=-1) # ---------------------------------------------------------------- shared middle joint block (state lever, round 2) # The weight inputs of the shared joint graph after (h, t_mod, gate), in input order args_3 .. args_32. Linear weights # are passed transposed ([1, in, out]), so each projection is one BATCH_MATMUL with a runtime operand. SHARED_WEIGHT_NAMES = ['repo_norm_w', 'repo_norm_b', 'repo_gate_t', 'repo_content_t', 'repo_ang_t', 'q_t', 'q_b', 'k_t', 'k_b', 'v_t', 'v_b', 'gq', 'gk', 'out_t', 'out_b', 'fc1_t', 'fc1_b', 'fc2_t', 'fc2_b', 'shift', 'lqkv_a_t', 'lq_b_t', 'lk_b_t', 'lv_b_t', 'lout_a_t', 'lout_b_t', 'lfc1_a_t', 'lfc1_b_t', 'lfc2_a_t', 'lfc2_b_t'] def shared_block_weights(jb): """A middle DenJoint's weights in SHARED_WEIGHT_NAMES order (float32 tensors, the shapes of the graph inputs).""" def tw(lin): return lin.weight.detach().float().t().contiguous()[None] def row(v): return v.detach().float()[None, None] b = jb.block r = b.repo return [row(r.norm_w), row(r.norm_b), tw(r.gate), tw(r.content), tw(r.ang), tw(b.q), row(b.q.bias), tw(b.k), row(b.k.bias), tw(b.v), row(b.v.bias), b.gq.detach().float(), b.gk.detach().float(), tw(b.out), row(b.out.bias), tw(b.fc1), row(b.fc1.bias), tw(b.fc2), row(b.fc2.bias), b.shift_table.detach().float(), tw(b.qkv_a), tw(b.q_b), tw(b.k_b), tw(b.v_b), tw(b.out_a), tw(b.out_b), tw(b.fc1_a), tw(b.fc1_b), tw(b.fc2_a), tw(b.fc2_b)] class DenJointShared(nn.Module): """One graph for every middle joint block: (h [1, L, 1024], t_mod [1, 6144], gate [1, 64], *weights) -> h, the block's weights as runtime inputs (shared_block_weights order). Same arithmetic as DenBlock.forward; compiled once per photo on the phone instead of once per block.""" def __init__(self, ref_joint): super().__init__() b = ref_joint.block self.c, self.h, self.d, self.chunk = b.c, b.h, b.d, b.chunk self.eps1, self.eps2, self.repo_eps, self.repo_half = b.eps1, b.eps2, b.repo.eps, b.repo.half def forward(self, x, mod, gate, *w): W = dict(zip(SHARED_WEIGHT_NAMES, w)) c, H, d = self.c, self.h, self.d bsz, L, _ = x.shape mm = torch.matmul g = gate[:, None, :] hr = G.safe_layer_norm(x, self.repo_eps, W['repo_norm_w'], W['repo_norm_b']) a = mm(F.silu(mm(hr, W['repo_gate_t'])) * mm(hr, W['repo_content_t']), W['repo_ang_t']) a = a.reshape(bsz, L, H, self.repo_half) cos, sin = torch.cos(a), torch.sin(a) m = mod + W['shift'] shift_msa, scale_msa, gate_msa = m[:, 0:c], m[:, c:2 * c], m[:, 2 * c:3 * c] shift_mlp, scale_mlp, gate_mlp = m[:, 3 * c:4 * c], m[:, 4 * c:5 * c], m[:, 5 * c:6 * c] hh = G.safe_layer_norm(x, self.eps1) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) half = d // 2 def rope(t): t0, t1 = t[..., :half], t[..., half:] return torch.cat([t0 * cos - t1 * sin, t0 * sin + t1 * cos], dim=-1) ax = mm(hh, W['lqkv_a_t']) * g q = G.safe_rms_unit(rope((mm(hh, W['q_t']) + W['q_b'] + mm(ax, W['lq_b_t'])).reshape(bsz, L, H, d))) * W['gq'] k = G.safe_rms_unit(rope((mm(hh, W['k_t']) + W['k_b'] + mm(ax, W['lk_b_t'])).reshape(bsz, L, H, d))) * W['gk'] v = (mm(hh, W['v_t']) + W['v_b'] + mm(ax, W['lv_b_t'])).reshape(bsz, L, H, d) o = G.attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.chunk) o = o.reshape(bsz, L, c) x = x + (mm(o, W['out_t']) + W['out_b'] + mm(mm(o, W['lout_a_t']) * g, W['lout_b_t'])) * gate_msa.unsqueeze(1) hh = G.safe_layer_norm(x, self.eps2) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) u = G._gelu_tanh(mm(hh, W['fc1_t']) + W['fc1_b'] + mm(mm(hh, W['lfc1_a_t']) * g, W['lfc1_b_t'])) return x + (mm(u, W['fc2_t']) + W['fc2_b'] + mm(mm(u, W['lfc2_a_t']) * g, W['lfc2_b_t'])) * gate_mlp.unsqueeze(1) # ---------------------------------------------------------------- shared middle Gaussian decoder graph (state lever) # One GaussianBlockGPU's weights as graph inputs, in input order: self-attention, cross-attention, norm2's affine, MLP. # Linear weights are passed transposed ([1, in, out]), as for the shared joint block. GAUSS_BLOCK_WEIGHT_NAMES = ['sa_q_t', 'sa_q_b', 'sa_k_t', 'sa_k_b', 'sa_v_t', 'sa_v_b', 'sa_gq', 'sa_gk', 'sa_out_t', 'sa_out_b', 'ca_q_t', 'ca_q_b', 'ca_k_t', 'ca_k_b', 'ca_v_t', 'ca_v_b', 'ca_gq', 'ca_gk', 'ca_out_t', 'ca_out_b', 'n2_w', 'n2_b', 'fc1_t', 'fc1_b', 'fc2_t', 'fc2_b'] def gauss_block_weights(gb): """One GaussianBlockGPU's weights in GAUSS_BLOCK_WEIGHT_NAMES order (float32 tensors, the graph input shapes).""" def tw(lin): return lin.weight.detach().float().t().contiguous()[None] def row(v): return v.detach().float().reshape(1, 1, -1) def mha(m): return [tw(m.q), row(m.q.bias), tw(m.k), row(m.k.bias), tw(m.v), row(m.v.bias), m.gq.detach().float(), m.gk.detach().float(), tw(m.out), row(m.out.bias)] assert gb.n2w is not None and gb.n2b is not None return (mha(gb.self_attn) + mha(gb.cross_attn) + [row(gb.n2w), row(gb.n2b), tw(gb.fc1), row(gb.fc1.bias), tw(gb.fc2), row(gb.fc2.bias)]) def gauss_shared_weights(g): """A middle Gauss graph's weights: its blocks' gauss_block_weights one after another.""" return [t for b in g.blocks for t in gauss_block_weights(b)] class GaussShared(nn.Module): """One graph for the Gaussian decoder's middle graphs (gs_1 .. gs_6, two blocks each): (h [1, A, 1024], latent [1, 8192, 16], *weights) -> h, the blocks' weights as runtime inputs (gauss_shared_weights order). Same arithmetic as Gauss.forward for a middle graph; compiled once per photo on the phone instead of six times.""" def __init__(self, ref): super().__init__() assert not ref.first and not ref.last b0 = ref.blocks[0] self.n = len(ref.blocks) self.eps = (b0.e1, b0.e2, b0.e3) self.sa = (b0.self_attn.h, b0.self_attn.c, b0.self_attn.d, b0.self_attn.chunk) self.ca = (b0.cross_attn.h, b0.cross_attn.c, b0.cross_attn.d, b0.cross_attn.chunk) @staticmethod def _mha(x, src, W, p, heads, c, d, chunk): mm = torch.matmul b, lq, _ = x.shape lk = src.shape[1] q = G.safe_rms_unit((mm(x, W[p + 'q_t']) + W[p + 'q_b']).reshape(b, lq, heads, d)) * W[p + 'gq'] k = G.safe_rms_unit((mm(src, W[p + 'k_t']) + W[p + 'k_b']).reshape(b, lk, heads, d)) * W[p + 'gk'] v = (mm(src, W[p + 'v_t']) + W[p + 'v_b']).reshape(b, lk, heads, d) o = G.attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), chunk) return mm(o.reshape(b, lq, c), W[p + 'out_t']) + W[p + 'out_b'] def forward(self, h, latent, *w): n = len(GAUSS_BLOCK_WEIGHT_NAMES) assert len(w) == n * self.n, len(w) e1, e2, e3 = self.eps for i in range(self.n): W = dict(zip(GAUSS_BLOCK_WEIGHT_NAMES, w[i * n:(i + 1) * n])) x = G.safe_layer_norm(h, e1) h = h + self._mha(x, x, W, 'sa_', *self.sa) h = h + self._mha(G.safe_layer_norm(h, e2, W['n2_w'], W['n2_b']), latent, W, 'ca_', *self.ca) u = G._gelu_tanh(torch.matmul(G.safe_layer_norm(h, e3), W['fc1_t']) + W['fc1_b']) h = h + (torch.matmul(u, W['fc2_t']) + W['fc2_b']) return h