Blame for src/model.py

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


def zero_out(blk):
    """Обнулить выходы внимания: новый блок на старте тождественен (FFN-выход и так zero-init)."""
    for name in ("o", "xo"):
        lin = getattr(blk, name, None)
        if isinstance(lin, nn.Linear):
            nn.init.zeros_(lin.weight)
            if lin.bias is not None:
                nn.init.zeros_(lin.bias)
    return blk


GRAFT_DEFAULT = {"e0x": 1, "d0x": 1, "e1": 1, "d1": 1, "e2": 2, "d2": 2, "e3": 2}


def sincos_2d(gh: int, gw: int, dim: int, device, dtype=None):
    """Абсолютные 2D-позиции (DiT-стиль sin/cos по каждой оси) — в паре с RoPE в attention."""
    key = ("pe2d", gh, gw, dim, str(device))
    tab = _ROPE_CACHE.get(key)
    if tab is None:
        half = dim // 2
        f = 1.0 / (10000 ** (torch.arange(half // 2, device=device, dtype=torch.float32) / (half / 2)))
        y = torch.arange(gh, device=device, dtype=torch.float32)[:, None] * f[None]
        x = torch.arange(gw, device=device, dtype=torch.float32)[:, None] * f[None]
        ey = torch.cat([y.sin(), y.cos()], -1)[:, None, :].expand(gh, gw, half)
        ex = torch.cat([x.sin(), x.cos()], -1)[None, :, :].expand(gh, gw, half)
        tab = torch.cat([ey, ex], -1).reshape(gh * gw, dim)
        _ROPE_CACHE[key] = tab
    return tab if dtype is None else tab.to(dtype)


class Body(nn.Module):
    """U-стек: ширина 192 → 512 → 1024 → 1536 с аддитивными скипами.

    Базовая архитектура = 36 блоков: L0 10 conv + 4 attn, L1 9, L2 9, L3 4 (238.5M, 44.4 GFLOP на 256 px).
    Привитые блоки входят в GRAFT_DEFAULT и на конструкторе тождественны — так старые чекпойнты
    (26 блоков) грузятся без сломки: лишние блоки остаются тождественными.
    """

    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
        self.tdim = text_dim if xatt else None
        heads = {0: 6, 1: 8, 2: 8, 3: 8}
        self.stem = nn.Linear(latent_ch, w0)
        self.pe_dim = 64
        self.pe = nn.Linear(self.pe_dim, w0, bias=False)      # zero-init: на старте не вмешивается
        nn.init.zeros_(self.pe.weight)
        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.e0x = nn.ModuleList()     # прививаемые L0-attn (пусто = архитектура без изменений)
        self.d0x = nn.ModuleList()
        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)
        self.graft(GRAFT_DEFAULT)     # базовая архитектура уже с прививкой (блоки тождественны)

    # ---- прививка блоков
    def graft(self, plan: dict) -> int:
        """Добавить блоки к уже обученной модели так, чтобы они стартовали тождеством.

        plan: {"e1": 2, "d1": 2, "e3": 2}; ключи — имена списков блоков.
        Тождество: FFN-выход (`down`) zero-init, выходы внимания (`o`, `xo`) обнуляем здесь.
        """
        heads = {0: 6, 1: 8, 2: 8, 3: 8}
        lvl = {"e0": 0, "d0": 0, "e0x": 0, "d0x": 0, "e1": 1, "d1": 1, "e2": 2, "d2": 2, "e3": 3}
        dev = next(self.parameters()).device
        added = 0
        for key, n in plan.items():
            seq = getattr(self, key)
            for _ in range(int(n)):
                if key in ("e0", "d0"):
                    blk = ConvBlock(self.widths[0], text_dim=self.tdim)
                else:
                    blk = AttnBlock(self.widths[lvl[key]], heads[lvl[key]], self.tdim)
                seq.append(zero_out(blk).to(dev))
                added += 1
        return added

    # ---- переходы
    @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
        # pooled-текст в adaLN: среднее ТОЛЬКО по реальным токенам (паддинг иначе разбавляет вектор,
        # и поведение модели начинает зависеть от длины паддинга, а не от подписи)
        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))
        x = x + self.pe(sincos_2d(gh0, gw0, self.pe_dim, x.device, x.dtype))
        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)
        for blk in self.e0x:
            x = blk(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.d0x:
            x = blk(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)