Blame for vae/ae/model.py

"""Image autoencoder: deterministic transformer encoder + flow-matching decoder.

* Encoder: 8x8 patches -> windowed-attention transformer blocks -> 32-channel latent
  on a grid 8x smaller than the image.
* Decoder: a conditional flow-matching model. Given the latent, a noisy image ``x_t`` and
  the time ``t``, it predicts the clean image ``x_0``. Images are produced by integrating
  from pure noise (t=1) to the image (t=0) in a small number of steps (default 4).

Everything except the decoder's upsampling head works on token sequences, so a batch may
mix images of any sizes and aspect ratios: each image is cut into its own grid of patch
tokens, the grids are packed into one padded sequence batch (:class:`Packing`), and
attention is restricted to each image's own grid. The convolutional upsampling head runs
per group of equal-shape images.

All architecture hyper-parameters come from :class:`AEConfig`.
"""
import json
import math
import os
from dataclasses import asdict, dataclass
from typing import List, Optional, Sequence, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from .window_attention import _tile_width, packed_window_attention


@dataclass
class AEConfig:
    in_channels: int = 3
    patch_size: int = 8               # spatial compression factor
    latent_channels: int = 32
    # encoder
    enc_width: int = 1024
    enc_depth: int = 8
    enc_mlp_ratio: float = 2.0
    # decoder
    dec_width: int = 512
    dec_depth: int = 6
    dec_mlp_ratio: float = 4.0
    dec_adaln_rank: int = 128         # rank of the per-block low-rank modulation deltas
    dec_skip_blocks: int = 2          # long skip: output of block (n-1) feeds the last n blocks
    dec_head_min_channels: int = 64
    time_embed_dim: int = 256
    # shared
    head_dim: int = 64
    window_radius: int = 8            # attention window: |dy| <= r and |dx| <= r tokens
    rope_base: float = 100.0
    # per-channel statistics of raw latents (used by encode()/decode())
    latent_mean: Optional[List[float]] = None
    latent_std: Optional[List[float]] = None

    @classmethod
    def from_json(cls, path):
        with open(path) as f:
            d = json.load(f)
        known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__}
        return cls(**known)

    def to_json(self, path):
        with open(path, "w") as f:
            json.dump(asdict(self), f, indent=2)


# --------------------------------------------------------------------------- packing

class Packing:
    """Layout of a batch of token grids of sizes ``hw = [(h, w), ...]``.

    Sample ``b`` is stored row-major in its own padded ``gyp x gxp`` grid at the start of a
    sequence of length ``N`` (the largest padded area); its real tokens are the top-left
    ``h x w`` corner. ``pos[b]`` are the sequence indices of its real tokens.
    """

    def __init__(self, hw: Sequence[Tuple[int, int]], radius: int, rope_base: float,
                 head_dim: int, device):
        self.hw = [(int(h), int(w)) for h, w in hw]
        WY, WX = 8, _tile_width(radius)
        B = len(self.hw)
        geo = [(math.ceil(h / WY) * WY, math.ceil(w / WX) * WX) for h, w in self.hw]
        self.N = max(a * b for a, b in geo)
        self.max_tiles = max((a // WY) * (b // WX) for a, b in geo)
        self.dims = torch.tensor([[a, b, h, w] for (a, b), (h, w) in zip(geo, self.hw)],
                                 dtype=torch.int32, device=device)
        self.valid = torch.zeros(B, self.N, 1, device=device)
        pos_y = torch.zeros(B, self.N, device=device)
        pos_x = torch.zeros(B, self.N, device=device)
        self.pos = []
        for b, ((_, gxp), (h, w)) in enumerate(zip(geo, self.hw)):
            yy = torch.arange(h, device=device).repeat_interleave(w)
            xx = torch.arange(w, device=device).repeat(h)
            p = yy * gxp + xx
            self.pos.append(p)
            self.valid[b, p] = 1
            pos_y[b, p], pos_x[b, p] = yy.float(), xx.float()
        # 2D axial RoPE tables: the head dim is split into four quarters,
        # two rotate with the x coordinate, two with the y coordinate.
        dq = head_dim // 4
        fr = torch.exp(-math.log(rope_base) * torch.arange(dq, device=device, dtype=torch.float32) / dq)
        ax, ay = pos_x[..., None] * fr, pos_y[..., None] * fr                  # (B, N, dq)
        self._rope = [t[:, None] for t in (ax.cos(), ax.sin(), ay.cos(), ay.sin())]
        self._dq = dq
        self.radius = radius

    def pack(self, seqs):
        """list of (n_b, C) -> (B, N, C), zero padded."""
        out = seqs[0].new_zeros(len(seqs), self.N, seqs[0].shape[-1])
        for b, s in enumerate(seqs):
            out[b, self.pos[b]] = s
        return out

    def unpack(self, x):
        """(B, N, C) -> list of (C, h, w)."""
        return [x[b, self.pos[b]].transpose(0, 1).reshape(-1, h, w) for b, (h, w) in enumerate(self.hw)]

    def rope(self, q, k):
        dq = self._dq
        cx, sx, cy, sy = (c.to(q.dtype) for c in self._rope)

        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)

    def attention(self, q, k, v):
        return packed_window_attention(q, k, v, self.radius, self.dims, self.max_tiles)


def patchify(img, p):
    """(C, H, W) -> (H/p * W/p, C*p*p) row-major patch tokens."""
    C, H, W = img.shape
    return img.reshape(C, H // p, p, W // p, p).permute(1, 3, 0, 2, 4).reshape(-1, C * p * p)


# --------------------------------------------------------------------------- encoder

class RMSNorm(nn.Module):
    def __init__(self, ch, eps=1e-8):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(ch))

    def forward(self, x):
        return x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() * self.weight


class Mlp(nn.Module):
    def __init__(self, i, h, o):
        super().__init__()
        self.m = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Linear(h, o))

    def forward(self, x):
        return self.m(x)


class PlainMlp(nn.Module):
    def __init__(self, d, ratio):
        super().__init__()
        h = int(d * ratio)
        self.up = nn.Linear(d, h)
        self.down = nn.Linear(h, d)
        self.act = nn.GELU(approximate="tanh")

    def forward(self, x):
        return self.down(self.act(self.up(x)))


class EncoderBlock(nn.Module):
    """Pre-norm transformer block with 2D-RoPE windowed self-attention and an MLP."""

    def __init__(self, ch, head_dim, mlp_ratio):
        super().__init__()
        self.h = max(1, ch // head_dim)
        self.norm1 = RMSNorm(ch)
        self.qkv = nn.Linear(ch, ch * 3, bias=False)
        self.proj = nn.Linear(ch, ch, bias=False)
        self.norm2 = RMSNorm(ch)
        self.mlp = PlainMlp(ch, mlp_ratio)

    def forward(self, x, pk: Packing):
        B, N, C = x.shape
        q, k, v = self.qkv(self.norm1(x)).reshape(B, N, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
        q, k = pk.rope(q, k)
        a = pk.attention(q, k, v).transpose(1, 2).reshape(B, N, C).to(x.dtype)
        x = x + self.proj(a)
        return x + self.mlp(self.norm2(x))


class Encoder(nn.Module):
    """patch tokens (B, N, 3*p*p) -> raw latent tokens (B, N, latent_channels)."""

    def __init__(self, cfg: AEConfig):
        super().__init__()
        w = cfg.enc_width
        self.initial = nn.Linear(cfg.in_channels * cfg.patch_size ** 2, w)
        self.blocks = nn.ModuleList(EncoderBlock(w, cfg.head_dim, cfg.enc_mlp_ratio)
                                    for _ in range(cfg.enc_depth))
        self.out = Mlp(w, w, cfg.latent_channels)

    def forward(self, x, pk: Packing):
        valid = pk.valid.to(x.dtype)
        x = self.initial(x)
        for b in self.blocks:
            x = b(x, pk) * valid
        return self.out(x)


# --------------------------------------------------------------------------- decoder

class DecoderBlock(nn.Module):
    """Transformer block with windowed 2D-RoPE attention, modulated by the timestep:
    the modulation vector carries (scale, gate) for the attention and the MLP sub-layers."""

    def __init__(self, ch, head_dim, mlp_ratio):
        super().__init__()
        self.h = max(1, ch // head_dim)
        self.norm1 = nn.RMSNorm(ch)
        self.qkv = nn.Linear(ch, 3 * ch)
        self.proj = nn.Linear(ch, ch)
        self.norm2 = nn.RMSNorm(ch)
        hid = int(ch * mlp_ratio)
        self.fc1 = nn.Linear(ch, hid)
        self.fc2 = nn.Linear(hid, ch)

    def forward(self, x, mod, pk: Packing):
        B, N, C = x.shape
        s1, g1, s2, g2 = mod[:, None, :].chunk(4, dim=-1)
        h = self.norm1(x) * (1 + s1)
        q, k, v = self.qkv(h).reshape(B, N, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
        q, k = pk.rope(q, k)
        a = pk.attention(q, k, v).transpose(1, 2).reshape(B, N, C).to(h.dtype)
        x = x + g1 * self.proj(a)
        h = self.norm2(x) * (1 + s2)
        return x + g2 * self.fc2(F.gelu(self.fc1(h)))


class FlowDecoder(nn.Module):
    """x0-prediction network: (x_t, t, latent) -> predicted clean image.

    The transformer works on packed tokens; the upsampling head (PixelShuffle + 3x3 convs)
    runs per group of equal-shape samples.
    """

    def __init__(self, cfg: AEConfig):
        super().__init__()
        D, p = cfg.dec_width, cfg.patch_size
        self.depth = cfg.dec_depth
        self.n_skip = cfg.dec_skip_blocks
        self.time_embed_dim = cfg.time_embed_dim
        self.patchify = nn.Linear(cfg.in_channels * p * p, D)
        self.latent_up = nn.Linear(cfg.latent_channels, D)
        self.fuse_in = nn.Linear(2 * D, D)
        self.tmlp = nn.Sequential(nn.Linear(cfg.time_embed_dim, D), nn.SiLU(), nn.Linear(D, D))
        self.adaln_act = nn.SiLU()
        mw = 4 * D
        self.adaln_base = nn.Linear(D, mw)
        self.d_down = nn.ModuleList(nn.Linear(D, cfg.dec_adaln_rank, bias=False) for _ in range(self.depth))
        self.d_up = nn.ModuleList(nn.Linear(cfg.dec_adaln_rank, mw, bias=False) for _ in range(self.depth))
        self.blocks = nn.ModuleList(DecoderBlock(D, cfg.head_dim, cfg.dec_mlp_ratio)
                                    for _ in range(self.depth))
        self.fuse_skip = nn.Linear(2 * D, D)
        # progressive upsampling head: [PixelShuffle(2) -> conv3x3 -> GELU] x log2(p)
        n = int(round(math.log2(p)))
        ch, stages = D, []
        for _ in range(n):
            nxt = max(ch // 2, cfg.dec_head_min_channels)
            stages.append(nn.Sequential(nn.PixelShuffle(2), nn.Conv2d(ch // 4, nxt, 3, padding=1), nn.GELU()))
            ch = nxt
        self.up_stages = nn.ModuleList(stages)
        self.up_final = nn.Conv2d(ch, cfg.in_channels, 3, padding=1)

    def time_embed(self, t):
        half = self.time_embed_dim // 2
        freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
        a = t.float()[:, None] * 1000.0 * freqs[None]
        return self.tmlp(torch.cat([a.sin(), a.cos()], dim=-1).to(self.tmlp[0].weight.dtype))

    def head(self, feats):
        """feats: (B, D, h, w) -> (B, 3, h*p, w*p)."""
        for st in self.up_stages:
            feats = st(feats)
        return self.up_final(feats)

    def forward(self, xt, t, z, pk: Packing):
        """xt: packed noisy patch tokens (B, N, 3*p*p); t: (B,); z: packed latents (B, N, C).
        Returns a list of per-sample x0 predictions (3, H_b, W_b)."""
        valid = pk.valid
        cond = self.adaln_act(self.time_embed(t))
        base = self.adaln_base(cond)
        x = self.fuse_in(torch.cat([self.patchify(xt), self.latent_up(z)], dim=-1))
        valid = valid.to(x.dtype)
        x = x * valid
        outs = []
        for i, b in enumerate(self.blocks):
            if i == self.depth - self.n_skip:
                x = self.fuse_skip(torch.cat([outs[self.n_skip - 1], x], dim=-1))
            x = b(x, base + self.d_up[i](self.d_down[i](cond)), pk) * valid
            outs.append(x)
        feats = pk.unpack(x)
        preds = [None] * len(feats)
        groups = {}
        for i, f in enumerate(feats):
            groups.setdefault(tuple(f.shape), []).append(i)
        for idx in groups.values():
            y = self.head(torch.stack([feats[i] for i in idx]))
            for j, i in enumerate(idx):
                preds[i] = y[j]
        return preds


# --------------------------------------------------------------------------- wrapper

def _as_list(x):
    """Tensor (B, C, H, W) or list of (C, H, W) -> (list, was_tensor)."""
    if torch.is_tensor(x):
        assert x.dim() == 4, f"expected (B, C, H, W), got {tuple(x.shape)}"
        return list(x.unbind(0)), True
    return list(x), False


class AutoEncoder(nn.Module):
    """Encoder + flow-matching decoder with built-in latent normalization.

    ``encode`` returns latents normalized per channel ((z - mean) / std), ready to be used
    as a diffusion target; ``decode`` expects such latents. ``encode_raw`` / ``decode_raw``
    skip the normalization.

    All four methods take either a batch tensor ``(B, C, H, W)`` (returns a tensor) or a
    list of ``(C, H_i, W_i)`` tensors of different sizes (returns a list), processed in one
    forward pass.
    """

    def __init__(self, cfg: AEConfig):
        super().__init__()
        self.config = cfg
        self.encoder = Encoder(cfg)
        self.decoder = FlowDecoder(cfg)
        C = cfg.latent_channels
        mean = torch.tensor(cfg.latent_mean, dtype=torch.float32) if cfg.latent_mean else torch.zeros(C)
        std = torch.tensor(cfg.latent_std, dtype=torch.float32) if cfg.latent_std else torch.ones(C)
        self.register_buffer("latent_mean", mean.view(C, 1, 1), persistent=False)
        self.register_buffer("latent_std", std.view(C, 1, 1), persistent=False)
        # compute dtype: encoder weights are stored in it; the decoder keeps fp32 weights and
        # runs under autocast with this dtype (when it is not float32).
        self.compute_dtype = torch.float32

    # ---------------------------------------------------------------- loading
    @classmethod
    def from_pretrained(cls, path_or_repo, device=None, dtype=None, **hub_kwargs):
        """Load from a local directory or a Hugging Face repo id containing
        ``config.json`` and ``model.safetensors``.

        dtype: compute dtype (default: bfloat16 on CUDA, float32 on CPU).
        """
        from safetensors.torch import load_file

        if os.path.isdir(path_or_repo):
            cfg_path = os.path.join(path_or_repo, "config.json")
            w_path = os.path.join(path_or_repo, "model.safetensors")
        else:
            from huggingface_hub import hf_hub_download
            cfg_path = hf_hub_download(path_or_repo, "config.json", **hub_kwargs)
            w_path = hf_hub_download(path_or_repo, "model.safetensors", **hub_kwargs)
        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        if dtype is None:
            dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        model = cls(AEConfig.from_json(cfg_path))
        model.load_state_dict(load_file(w_path), strict=True)
        model.to(device).eval().requires_grad_(False)
        model.set_dtype(dtype)
        return model

    def set_dtype(self, dtype):
        self.compute_dtype = dtype
        self.encoder.to(dtype)
        self.decoder.float()
        return self

    @property
    def device(self):
        return self.latent_mean.device

    def _autocast(self):
        return torch.autocast(self.device.type, dtype=self.compute_dtype,
                              enabled=self.compute_dtype != torch.float32)

    def _packing(self, hw):
        c = self.config
        return Packing(hw, c.window_radius, c.rope_base, c.head_dim, self.device)

    # ---------------------------------------------------------------- encoder side
    @torch.no_grad()
    def encode_raw(self, images):
        """images: (B, 3, H, W) or a list of (3, H_i, W_i); float in [-1, 1] or uint8 in
        [0, 255]; sides multiples of the patch size.
        Returns raw (unnormalized) float32 latents: (B, C, H/p, W/p) or a list of (C, H_i/p, W_i/p)."""
        p = self.config.patch_size
        imgs, was_tensor = _as_list(images)
        for im in imgs:
            assert im.dim() == 3 and im.shape[-2] % p == 0 and im.shape[-1] % p == 0, \
                f"each image must be (3, H, W) with H, W multiples of {p}, got {tuple(im.shape)}"
        pk = self._packing([(im.shape[-2] // p, im.shape[-1] // p) for im in imgs])
        toks = []
        for im in imgs:
            im = im.to(self.device)
            im = im.float() / 127.5 - 1.0 if im.dtype == torch.uint8 else im
            toks.append(patchify(im.to(self.compute_dtype), p))
        z = [t.float() for t in pk.unpack(self.encoder(pk.pack(toks), pk))]
        return torch.stack(z) if was_tensor else z

    @torch.no_grad()
    def encode(self, images):
        """Images -> normalized latents (zero mean / unit std per channel)."""
        z = self.encode_raw(images)
        if torch.is_tensor(z):
            return (z - self.latent_mean) / self.latent_std
        return [(t - self.latent_mean) / self.latent_std for t in z]

    # ---------------------------------------------------------------- decoder side
    @torch.no_grad()
    def decode_raw(self, latents, steps=4, generator=None, noise=None):
        """Raw latents -> images (approximately in [-1, 1], not clamped).

        latents: (B, C, h, w) or a list of (C, h_i, w_i).
        steps: number of flow-matching integration steps (default 4; 1 gives a fast,
               softer result, more steps add a bit of detail).
        generator: optional torch.Generator for the initial noise (on the model device).
        noise: optional explicit initial noise, (B, 3, H, W) or a list of (3, H_i, W_i).
        """
        p, C = self.config.patch_size, self.config.in_channels
        zs, was_tensor = _as_list(latents)
        zs = [z.to(self.device, torch.float32) for z in zs]
        pk = self._packing([tuple(z.shape[-2:]) for z in zs])
        zp = pk.pack([z.flatten(1).transpose(0, 1) for z in zs])
        if noise is not None:
            xs = [n.to(self.device, torch.float32) for n in _as_list(noise)[0]]
        elif was_tensor:
            B, _, h, w = latents.shape
            xs = list(torch.randn((B, C, h * p, w * p), generator=generator, device=self.device).unbind(0))
        else:
            xs = [torch.randn((C, z.shape[1] * p, z.shape[2] * p), generator=generator, device=self.device)
                  for z in zs]
        ts = torch.linspace(1.0, 0.0, steps + 1, device=self.device)
        B = len(zs)
        for i in range(steps):
            with self._autocast():
                xt = pk.pack([patchify(x.to(self.compute_dtype), p) for x in xs])
                preds = self.decoder(xt, ts[i].expand(B), zp, pk)
            # x0 prediction -> implied noise -> move along the straight path to ts[i+1]
            xs = [ts[i + 1] * (x - (1 - ts[i]) * pr.float()) / ts[i] + (1 - ts[i + 1]) * pr.float()
                  for x, pr in zip(xs, preds)]
        return torch.stack(xs) if was_tensor else xs

    @torch.no_grad()
    def decode(self, latents, steps=4, generator=None, noise=None):
        """Normalized latents (as returned by ``encode``) -> images in about [-1, 1]."""
        if torch.is_tensor(latents):
            z = latents.to(self.device, torch.float32) * self.latent_std + self.latent_mean
        else:
            z = [t.to(self.device, torch.float32) * self.latent_std + self.latent_mean for t in latents]
        return self.decode_raw(z, steps=steps, generator=generator, noise=noise)

    @torch.no_grad()
    def forward(self, images, steps=4, generator=None):
        """Round trip: images -> latents -> images."""
        return self.decode(self.encode(images), steps=steps, generator=generator)