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