"""UDiT — пирамидальный U-трансформер: тело (латентный поток) + хвост (декодер AE). Идея: ширина растёт там, где токенов мало. Узко на высоком разрешении, широко внизу, как в U-Net'ах SD (320/640/1280 при 64/32/16), а не 1024 на всех уровнях. 512², латент f8 (64×36 = 2304 токена): L0 ÷1 2304 tok C=192 conv-блоки (depthwise 3×3 + expand 4×) 6 enc / 4 dec L1 ÷2 576 tok C=512 трансформер + cross-attn к тексту 4 enc / 3 dec L2 ÷4 144 tok C=1024 трансформер, полное внимание 3 enc / 2 dec L3 ÷8 36 tok C=1536 трансформер, полное внимание 2 (боттлнек) Скипы — аддитивные (SD-стиль): после e0/e1/e2 признаки прибавляются в декодере. Переходы: mean-pool 2× + Linear вниз, Linear + nearest 2× вверх. Оценка: ~160M параметров и ~70 GFLOP против ~580 GFLOP у плоского 1024-стека. Хвост — декодер AE (рендер латента в пиксели), текст — Ettin + SliceMixer (из `text.py`). """ from __future__ import annotations import math import sys from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT / "vae")) from ae.model import AutoEncoder # noqa: E402 DEV = "cuda" if torch.cuda.is_available() else "cpu" _ROPE_CACHE: dict = {} def t_embed(t: torch.Tensor, dim: int = 256) -> torch.Tensor: """Синусоидальный эмбеддинг времени (t в [0,1] умножается на 1000).""" half = dim // 2 f = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device, dtype=torch.float32) / half) a = (t.float()[:, None] * 1000.0) * f[None] return torch.cat([a.sin(), a.cos()], -1) def body_rope(q, k, gh, gw, base: float = 100.0): """2D RoPE: 4 четверти head_dim, две крутятся по x, две по y. Фаза — в латентных пикселях.""" key = (gh, gw, q.shape[-1], q.device, base) tabs = _ROPE_CACHE.get(key) if tabs is None: dq = q.shape[-1] // 4 fr = torch.exp(-math.log(base) * torch.arange(dq, device=q.device, dtype=torch.float32) / dq) yy = torch.arange(gh, device=q.device).repeat_interleave(gw).float() xx = torch.arange(gw, device=q.device).repeat(gh).float() ax, ay = xx[:, None] * fr, yy[:, None] * fr tabs = (ax.cos()[None, None], ax.sin()[None, None], ay.cos()[None, None], ay.sin()[None, None]) _ROPE_CACHE[key] = tabs cx, sx, cy, sy = tabs dq = q.shape[-1] // 4 def rot(t): a, b, c, e = t[..., :dq], t[..., dq:2 * dq], t[..., 2 * dq:3 * dq], t[..., 3 * dq:] return torch.cat([a * cx - b * sx, a * sx + b * cx, c * cy - e * sy, c * sy + e * cy], dim=-1) return rot(q), rot(k) class RMSNorm(nn.Module): def __init__(self, d: int, eps: float = 1e-6): super().__init__() self.d, self.eps = d, eps self.scale = nn.Parameter(torch.zeros(d)) def forward(self, x): return F.rms_norm(x.float(), (self.d,), weight=(self.scale.float() + 1.0), eps=self.eps).to(x.dtype) class ConvBlock(nn.Module): """ConvNeXt-ish: RMSNorm → depthwise 3×3 → expand 4× → GELU → project, всё под adaLN-гейтом.""" def __init__(self, c: int, expand: int = 4, text_dim: int | None = None, heads: int = 6): super().__init__() self.c, self.hid, self.h, self.hd = c, c * expand, heads, c // heads self.norm = RMSNorm(c) self.norm_x = RMSNorm(c) self.dw = nn.Conv2d(c, c, 3, padding=1, groups=c) self.up = nn.Linear(c, self.hid) self.down = nn.Linear(self.hid, c) self.mod = nn.Parameter(torch.zeros(5, c)) # shift, scale, gate | shift_x, scale_x self.xatt = text_dim is not None if self.xatt: # текст на мелком уровне: не экономим self.xq, self.xo = nn.Linear(c, c, bias=False), nn.Linear(c, c, bias=False) nn.init.zeros_(self.down.weight); nn.init.zeros_(self.down.bias) def forward(self, x, vec, gh, gw, text=None, text_mask=None): B = x.shape[0] shift, scale, gate, shift_x, scale_x = (self.mod + vec).unbind(1) h = self.norm(x) * (1 + scale[:, None]) + shift[:, None] h = self.dw(h.transpose(1, 2).reshape(B, self.c, gh, gw)).flatten(2).transpose(1, 2) x = x + gate[:, None] * self.down(F.gelu(self.up(h))) if self.xatt and text is not None: N = x.shape[1] qq = self.xq(self.norm_x(x) * (1 + scale_x[:, None]) + shift_x[:, None]) qq = qq.reshape(B, N, self.h, self.hd).transpose(1, 2) kk, vv = text mask = text_mask[:, None, None, :] if text_mask is not None else None xa = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask) x = x + self.xo(xa.transpose(1, 2).reshape(B, N, self.c)) return x class AttnBlock(nn.Module): """Pre-norm трансформер: self-attention с 2D RoPE + (опц.) cross-attn к тексту + FFN.""" def __init__(self, c: int, heads: int, xatt_dim: int | None, mlp: float = 2.0): super().__init__() self.c, self.h, self.hd = c, heads, c // heads self.n1, self.n2, self.n3 = RMSNorm(c), RMSNorm(c), RMSNorm(c) self.qkv = nn.Linear(c, 3 * c, bias=False) self.o = nn.Linear(c, c, bias=False) self.up, self.down = nn.Linear(c, int(c * mlp)), nn.Linear(int(c * mlp), c) self.mod = nn.Parameter(torch.zeros(6, c)) self.xatt = xatt_dim is not None if self.xatt: self.xq, self.xo = nn.Linear(c, c, bias=False), nn.Linear(c, c, bias=False) nn.init.zeros_(self.down.weight); nn.init.zeros_(self.down.bias) def forward(self, x, vec, gh, gw, text=None, text_mask=None): B, N, C = x.shape s1, g1, s2, g2, s3, g3 = (self.mod + vec).unbind(1) h = self.n1(x) * (1 + s1[:, None]) q, k, v = self.qkv(h).reshape(B, N, 3, self.h, self.hd).permute(2, 0, 3, 1, 4) q, k = body_rope(q, k, gh, gw) a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, C) x = x + g1[:, None] * self.o(a) if self.xatt and text is not None: qq = self.xq(self.n2(x)).reshape(B, N, self.h, self.hd).transpose(1, 2) kk, vv = text mask = text_mask[:, None, None, :] if text_mask is not None else None xa = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask) x = x + g2[:, None] * self.xo(xa.transpose(1, 2).reshape(B, N, C)) h = self.n3(x) * (1 + s3[:, None]) return x + g3[:, None] * self.down(F.gelu(self.up(h))) class Body(nn.Module): """U-стек: ширина 192 → 512 → 1024 → 1536 с аддитивными скипами.""" def __init__(self, latent_ch: int = 32, text_dim: int = 768, widths: tuple = (192, 512, 1024, 1536), enc: tuple = (6, 4, 3, 2), dec: tuple = (4, 3, 2), xatt: bool = True): super().__init__() w0, w1, w2, w3 = widths self.widths, self.enc, self.dec = widths, enc, dec heads = {0: 6, 1: 8, 2: 8, 3: 8} self.stem = nn.Linear(latent_ch, w0) self.e0 = nn.ModuleList([ConvBlock(w0, text_dim=text_dim if xatt else None) for _ in range(enc[0])]) self.e0a = AttnBlock(w0, 6, text_dim if xatt else None) self.d0a = AttnBlock(w0, 6, text_dim if xatt else None) self.e1 = nn.ModuleList([AttnBlock(w1, heads[1], text_dim if xatt else None) for _ in range(enc[1])]) self.e2 = nn.ModuleList([AttnBlock(w2, heads[2], text_dim if xatt else None) for _ in range(enc[2])]) self.e3 = nn.ModuleList([AttnBlock(w3, heads[3], text_dim if xatt else None) for _ in range(enc[3])]) self.d2 = nn.ModuleList([AttnBlock(w2, heads[2], text_dim if xatt else None) for _ in range(dec[2])]) self.d1 = nn.ModuleList([AttnBlock(w1, heads[1], text_dim if xatt else None) for _ in range(dec[1])]) self.d0 = nn.ModuleList([ConvBlock(w0, text_dim=text_dim if xatt else None) for _ in range(dec[0])]) self.dn01, self.dn12, self.dn23 = nn.Linear(w0, w1), nn.Linear(w1, w2), nn.Linear(w2, w3) self.up32, self.up21, self.up10 = nn.Linear(w3, w2), nn.Linear(w2, w1), nn.Linear(w1, w0) self.tmlp = nn.Sequential(nn.Linear(256, 512), nn.SiLU(), nn.Linear(512, 512)) self.pmlp = nn.Linear(text_dim, 512) self.ada = nn.ModuleDict({f"l{i}": nn.Sequential(nn.SiLU(), nn.Linear(512, (5 if i == 0 else 6) * w)) for i, w in enumerate(widths)}) self.ada["l0a"] = nn.Sequential(nn.SiLU(), nn.Linear(512, 6 * w0)) # для self-attn на L0 self.xk0 = nn.Linear(text_dim, 2 * w0) self.xk1, self.xk2, self.xk3 = (nn.Linear(text_dim, 2 * w1), nn.Linear(text_dim, 2 * w2), nn.Linear(text_dim, 2 * w3)) self.out_norm = RMSNorm(w0) self.out = nn.Linear(w0, latent_ch) nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias) # ---- переходы @staticmethod def downsample(x, gh, gw, proj): B = x.shape[0] d = x.view(B, gh // 2, 2, gw // 2, 2, -1).mean(dim=(2, 4)) return proj(d.reshape(B, (gh // 2) * (gw // 2), -1)) @staticmethod def upsample(x, gh, gw, proj): """(B, gh*gw, C) -> (B, 4*gh*gw, C') в родительском row-major порядке.""" B = x.shape[0] d = proj(x).view(B, gh, 1, gw, 1, -1).expand(B, gh, 2, gw, 2, proj.out_features) return d.reshape(B, 4 * gh * gw, proj.out_features) def _txt(self, proj, text, heads): B, L, _ = text.shape kk, vv = proj(text).chunk(2, dim=-1) hd = kk.shape[-1] // heads return (kk.reshape(B, L, heads, hd).transpose(1, 2), vv.reshape(B, L, heads, hd).transpose(1, 2)) def forward(self, z, t, text, text_mask, gh=None, gw=None): """gh/gw принимаются для совместимости с интерфейсом UDM и игнорируются.""" B, C, gh0, gw0 = z.shape cond = self.tmlp(t_embed(t, 256)) + self.pmlp(text.mean(1)) v = [self.ada[f"l{i}"](cond).view(B, 5 if i == 0 else 6, -1) for i in range(4)] v0a = self.ada["l0a"](cond).view(B, 6, -1) tk0 = self._txt(self.xk0, text, 6) tk1 = self._txt(self.xk1, text, 8) tk2 = self._txt(self.xk2, text, 8) tk3 = self._txt(self.xk3, text, 8) x = self.stem(z.flatten(2).transpose(1, 2)) for blk in self.e0: x = blk(x, v[0], gh0, gw0, tk0, text_mask) x = self.e0a(x, v0a, gh0, gw0, tk0, text_mask) s0 = x g1h, g1w = gh0 // 2, gw0 // 2 x = self.downsample(x, gh0, gw0, self.dn01) for blk in self.e1: x = blk(x, v[1], g1h, g1w, tk1, text_mask) s1 = x g2h, g2w = g1h // 2, g1w // 2 x = self.downsample(x, g1h, g1w, self.dn12) for blk in self.e2: x = blk(x, v[2], g2h, g2w, tk2, text_mask) s2 = x g3h, g3w = g2h // 2, g2w // 2 x = self.downsample(x, g2h, g2w, self.dn23) for blk in self.e3: x = blk(x, v[3], g3h, g3w, tk3, text_mask) # декодер: скипы аддитивные x = self.upsample(x, g3h, g3w, self.up32) + s2 for blk in self.d2: x = blk(x, v[2], g2h, g2w, tk2, text_mask) x = self.upsample(x, g2h, g2w, self.up21) + s1 for blk in self.d1: x = blk(x, v[1], g1h, g1w, tk1, text_mask) x = self.upsample(x, g1h, g1w, self.up10) + s0 x = self.d0a(x, v0a, gh0, gw0, tk0, text_mask) for blk in self.d0: x = blk(x, v[0], gh0, gw0, tk0, text_mask) return self.out(self.out_norm(x)).transpose(1, 2).reshape(B, C, gh0, gw0) class UDiT(nn.Module): """Пирамидальное тело + хвост-декодер AE; интерфейс как у UDM.""" def __init__(self, ae: AutoEncoder, text_dim: int = 768, **kw): super().__init__() self.ae = ae self.body = Body(latent_ch=ae.config.latent_channels, text_dim=text_dim, **kw) def to_raw(self, z): return z * self.ae.latent_std + self.ae.latent_mean @torch.no_grad() def sample(self, text, text_mask, gh, gw, steps=4, tail_steps=4, generator=None, param="x0", cfg=1.0, null_text=None, null_mask=None, clip_norm=0.0): dev = self.ae.device z = torch.randn(1, self.ae.config.latent_channels, gh, gw, device=dev, generator=generator) with torch.autocast(dev.type, torch.bfloat16, enabled=dev.type == "cuda"): for i in range(steps): t = torch.full((1,), 1.0 - i / steps, device=dev) o = self.body(z, t, text, text_mask).float() if cfg > 1.0 and null_text is not None: ou = self.body(z, t, null_text, null_mask).float() o = ou + cfg * (o - ou) v = o if param == "vel" else (z - o) / t.clamp(min=1e-2)[:, None, None, None] z = z - v / steps if clip_norm: z = z.clamp(-clip_norm, clip_norm) return self.ae.decode_raw(self.to_raw(z), steps=tail_steps, generator=generator)