simit 0.1.0__py3-none-any.whl

simit/backends/mot/kernels.py

@@ -0,0 +1,130 @@
+"""Fused Triton kernels for the MoT decoder's elementwise ops.
+
+Each kernel reproduces the eager PyTorch op sequence including its
+intermediate roundings to the activation dtype (one pass over memory instead
+of 3-7), so outputs match the eager modules up to fp32 reduction order.
+Without Triton (or on CPU) the eager versions are used.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import triton
+    import triton.language as tl
+    _HAVE_TRITON = True
+except Exception:  # pragma: no cover
+    _HAVE_TRITON = False
+
+
+def _enabled(x: torch.Tensor) -> bool:
+    return _HAVE_TRITON and x.is_cuda and x.dtype == torch.bfloat16
+
+
+if _HAVE_TRITON:
+    @triton.jit
+    def _bf16(x):
+        """Round fp32 to the nearest-even bfloat16, kept in fp32. Bit-level, because
+        the compiler folds a plain fp32->bf16->fp32 cast pair away."""
+        b = x.to(tl.uint32, bitcast=True)
+        b = b + (0x7FFF + ((b >> 16) & 1))
+        return ((b >> 16) << 16).to(tl.float32, bitcast=True)
+
+    @triton.jit
+    def _rms_norm_kernel(X, W, Y, M, N, eps, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
+        rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+        cols = tl.arange(0, BLOCK_N)
+        mask = (rows[:, None] < M) & (cols[None, :] < N)
+        offs = rows[:, None].to(tl.int64) * N + cols[None, :]
+        x = tl.load(X + offs, mask=mask, other=0.0).to(tl.float32)
+        rstd = 1.0 / tl.sqrt(tl.sum(x * x, axis=1) / N + eps)
+        normed = _bf16(x * rstd[:, None])                                    # x.to(dtype)
+        w = tl.load(W + cols, mask=cols < N, other=0.0).to(tl.float32)
+        y = w[None, :] * normed
+        if Y.dtype.element_ty == tl.bfloat16:
+            y = _bf16(y)
+        tl.store(Y + offs, y.to(Y.dtype.element_ty), mask=mask)
+
+    @triton.jit
+    def _rope_kernel(X, COS, SIN, Y, M, H, D: tl.constexpr, BLOCK_M: tl.constexpr):
+        # X: [T*H, D] rows; COS/SIN: [T, D]; out = x*cos + rotate_half(x)*sin
+        rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+        half: tl.constexpr = D // 2
+        cols = tl.arange(0, half)
+        rmask = rows[:, None] < M
+        base = rows[:, None].to(tl.int64) * D
+        t = (rows // H)[:, None].to(tl.int64) * D
+        x1 = tl.load(X + base + cols[None, :], mask=rmask, other=0.0)
+        x2 = tl.load(X + base + half + cols[None, :], mask=rmask, other=0.0)
+        c1 = tl.load(COS + t + cols[None, :], mask=rmask, other=0.0)
+        c2 = tl.load(COS + t + half + cols[None, :], mask=rmask, other=0.0)
+        s1 = tl.load(SIN + t + cols[None, :], mask=rmask, other=0.0)
+        s2 = tl.load(SIN + t + half + cols[None, :], mask=rmask, other=0.0)
+        a1 = _bf16(x1.to(tl.float32) * c1.to(tl.float32))
+        a2 = _bf16(x2.to(tl.float32) * c2.to(tl.float32))
+        b1 = _bf16(-x2.to(tl.float32) * s1.to(tl.float32))
+        b2 = _bf16(x1.to(tl.float32) * s2.to(tl.float32))
+        tl.store(Y + base + cols[None, :], _bf16(a1 + b1).to(tl.bfloat16), mask=rmask)
+        tl.store(Y + base + half + cols[None, :], _bf16(a2 + b2).to(tl.bfloat16), mask=rmask)
+
+    @triton.jit
+    def _silu_mul_kernel(G, U, Y, n, BLOCK: tl.constexpr):
+        offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < n
+        g = tl.load(G + offs, mask=mask, other=0.0).to(tl.float32)
+        u = tl.load(U + offs, mask=mask, other=0.0).to(tl.float32)
+        s = _bf16(g / (1.0 + tl.exp(-g)))                                     # F.silu(g)
+        tl.store(Y + offs, _bf16(s * u).to(tl.bfloat16), mask=mask)
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """``weight * (x.float() * rsqrt(mean(x^2) + eps)).to(x.dtype)``"""
+    if not _enabled(x) or x.numel() == 0:
+        dtype = x.dtype
+        xf = x.float()
+        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+        return weight * xf.to(dtype)
+    N = x.shape[-1]
+    x2 = x.contiguous().view(-1, N)
+    y = torch.empty(x.shape, dtype=torch.promote_types(weight.dtype, x.dtype), device=x.device)
+    M = x2.shape[0]
+    BLOCK_N = triton.next_power_of_2(N)
+    BLOCK_M = max(1, min(64, 4096 // BLOCK_N))
+    _rms_norm_kernel[(triton.cdiv(M, BLOCK_M),)](x2, weight, y, M, N, eps, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
+    return y
+
+
+def _rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """``x * cos + rotate_half(x) * sin`` for x [T, H, D] and cos/sin [T, D]."""
+    D = x.shape[-1]
+    if (not _enabled(x) or x.numel() == 0 or cos.dtype != x.dtype or D & (D - 1)
+            or cos.shape != (x.shape[0], D)):
+        c, s = cos[:, None], sin[:, None]
+        return x * c + _rotate_half(x) * s
+    T, H = x.shape[0], x.shape[1]
+    x = x.contiguous()
+    cos, sin = cos.contiguous(), sin.contiguous()
+    y = torch.empty_like(x)
+    M = T * H
+    BLOCK_M = 32
+    _rope_kernel[(triton.cdiv(M, BLOCK_M),)](x, cos, sin, y, M, H, D=D, BLOCK_M=BLOCK_M)
+    return y
+
+
+def silu_mul(g: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
+    """``F.silu(g) * u``"""
+    if not _enabled(g) or g.numel() == 0 or g.dtype != u.dtype:
+        return F.silu(g) * u
+    g, u = g.contiguous(), u.contiguous()
+    y = torch.empty_like(g)
+    n = g.numel()
+    BLOCK = 2048
+    _silu_mul_kernel[(triton.cdiv(n, BLOCK),)](g, u, y, n, BLOCK=BLOCK)
+    return y

simit/backends/mot/modeling.py

@@ -0,0 +1,514 @@
+"""Inference-only modules for BAGEL-style Mixture-of-Transformers UMMs.
+
+Re-implemented for batched serving on PyTorch's built-in variable-length
+attention (no flash-attn install needed). Parameter names match the
+released BAGEL checkpoint (``language_model.model.layers.*`` ...), which
+Lance reuses for its language model.
+
+The language model processes one flat batch of tokens per step. Tokens are
+grouped into "pseudo-sequences" (a decode token, a prefill chunk, an image
+block) that attend to a contiguous region of a shared KV arena; causal and
+bidirectional pseudo-sequences are attended in two calls per layer.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.attention.varlen import varlen_attn
+
+from .kernels import apply_rope, rms_norm, silu_mul
+
+
+class DeviceMixin:
+    """``.to(device)`` for model containers that hold several modules (load on
+    CPU once, move to the GPU per request -- e.g. on Hugging Face ZeroGPU)."""
+
+    def to(self, device):
+        device = torch.device(device)
+        for name, value in list(vars(self).items()):
+            if isinstance(value, nn.Module):
+                value.to(device)
+            elif isinstance(value, torch.Tensor):
+                setattr(self, name, value.to(device))
+        self.device = device
+        return self
+
+
+# ----------------------------------------------------------------- basic layers
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return rms_norm(x, self.weight, self.eps)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden: int, inter: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, inter, bias=False)
+        self.up_proj = nn.Linear(hidden, inter, bias=False)
+        self.down_proj = nn.Linear(inter, hidden, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(silu_mul(self.gate_proj(x), self.up_proj(x)))
+
+
+class RotaryEmbedding(nn.Module):
+    """1-D RoPE with cached cos/sin tables (positions are integers)."""
+
+    def __init__(self, head_dim: int, theta: float, max_positions: int = 65536):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.max_positions = 0
+        self._cos = self._sin = None
+        self._grow(max_positions)
+
+    def _grow(self, n: int):
+        t = torch.arange(n, dtype=torch.float32, device=self.inv_freq.device)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self._cos, self._sin = emb.cos(), emb.sin()
+        self.max_positions = n
+
+    def forward(self, positions: torch.Tensor, dtype):
+        if self._cos.device != positions.device:
+            self._cos, self._sin = self._cos.to(positions.device), self._sin.to(positions.device)
+        if int(positions.max()) >= self.max_positions:
+            self._grow(int(positions.max()) * 2)
+            self._cos, self._sin = self._cos.to(positions.device), self._sin.to(positions.device)
+        return self._cos[positions].to(dtype), self._sin[positions].to(dtype)
+
+
+class MRotaryEmbedding(RotaryEmbedding):
+    """Qwen2-VL multimodal RoPE: positions are (t, h, w) triples and the head
+    dimension is split into ``sections`` (e.g. [16, 24, 24]) per axis."""
+
+    def __init__(self, head_dim: int, theta: float, sections, max_positions: int = 65536):
+        super().__init__(head_dim, theta, max_positions)
+        axis = []
+        for i, n in enumerate(list(sections) * 2):
+            axis += [i % 3] * n
+        assert len(axis) == head_dim, (sections, head_dim)
+        self.register_buffer("axis_of_dim", torch.tensor(axis), persistent=False)
+
+    def forward(self, positions: torch.Tensor, dtype):
+        if positions.dim() == 1:
+            return super().forward(positions, dtype)
+        cos, sin = super().forward(positions.reshape(-1), dtype)       # [3*T, D]
+        T, D = positions.shape[1], cos.shape[-1]
+        cos, sin = cos.view(3, T, D), sin.view(3, T, D)
+        idx = self.axis_of_dim.to(positions.device)[None, None, :].expand(1, T, D)
+        return cos.gather(0, idx)[0], sin.gather(0, idx)[0]
+
+
+# --------------------------------------------------------------- batch metadata
+@dataclass
+class AttnBatch:
+    """Where every token of a step lives and what it may attend to.
+
+    Tokens are ordered causal pseudo-sequences first (``n_causal`` tokens),
+    then bidirectional ones. For each group: ``cu_q`` query offsets, ``k_start``
+    arena start of the region, ``k_used`` keys visible (prefix + chunk).
+    Generation-expert tokens (latents) come last: rows ``[n_und:]``, so the
+    experts run on contiguous slices instead of gathered rows."""
+
+    slots: torch.Tensor                 # [T] arena index each token's K/V is written to
+    n_causal: int
+    causal: Optional[tuple] = None      # (cu_q, k_start_cu, k_used, max_q, max_k)
+    full: Optional[tuple] = None
+    n_und: Optional[int] = None         # None -> all tokens use the understanding expert
+
+
+class KVArena:
+    """Per-layer contiguous K/V storage: ``k[layer]`` is [capacity, Hkv, D]."""
+
+    def __init__(self, num_layers: int, capacity: int, num_kv_heads: int, head_dim: int, dtype, device):
+        self.k = torch.empty(num_layers, capacity, num_kv_heads, head_dim, dtype=dtype, device=device)
+        self.v = torch.empty_like(self.k)
+        self.capacity = capacity
+
+    def copy(self, src: int, dst: int, n: int):
+        self.k[:, dst:dst + n].copy_(self.k[:, src:src + n])
+        self.v[:, dst:dst + n].copy_(self.v[:, src:src + n])
+
+
+def _attend(q, arena_k, arena_v, meta, window, scale):
+    cu_q, k_start, k_used, max_q, max_k = meta
+    return varlen_attn(q, arena_k, arena_v, cu_q, k_start, max_q, max_k, window_size=window,
+                       enable_gqa=q.shape[1] != arena_k.shape[1], seqused_k=k_used, scale=scale)
+
+
+# --------------------------------------------------------------- MoT decoder
+def _split(fn_und, fn_gen, x, n_und: Optional[int]):
+    """Rows [:n_und] through the understanding expert, the rest through the
+    generation expert (functions may return a tensor or a tuple of tensors)."""
+    if n_und is None or n_und == x.shape[0]:
+        return fn_und(x)
+    if n_und == 0:
+        return fn_gen(x)
+    a, b = fn_und(x[:n_und]), fn_gen(x[n_und:])
+    if isinstance(a, tuple):
+        return tuple(torch.cat(pair) for pair in zip(a, b))
+    return torch.cat([a, b])
+
+
+class MoTAttention(nn.Module):
+    def __init__(self, hidden, n_heads, n_kv, head_dim, eps, layer_idx):
+        super().__init__()
+        self.n_heads, self.n_kv, self.head_dim, self.layer_idx = n_heads, n_kv, head_dim, layer_idx
+        self.scale = head_dim ** -0.5
+        for suffix in ("", "_moe_gen"):
+            setattr(self, f"q_proj{suffix}", nn.Linear(hidden, n_heads * head_dim, bias=True))
+            setattr(self, f"k_proj{suffix}", nn.Linear(hidden, n_kv * head_dim, bias=True))
+            setattr(self, f"v_proj{suffix}", nn.Linear(hidden, n_kv * head_dim, bias=True))
+            setattr(self, f"o_proj{suffix}", nn.Linear(n_heads * head_dim, hidden, bias=False))
+            setattr(self, f"q_norm{suffix}", RMSNorm(head_dim, eps))
+            setattr(self, f"k_norm{suffix}", RMSNorm(head_dim, eps))
+
+    def _qkv(self, x, suffix):
+        q = getattr(self, f"q_proj{suffix}")(x).view(-1, self.n_heads, self.head_dim)
+        k = getattr(self, f"k_proj{suffix}")(x).view(-1, self.n_kv, self.head_dim)
+        v = getattr(self, f"v_proj{suffix}")(x).view(-1, self.n_kv, self.head_dim)
+        return getattr(self, f"q_norm{suffix}")(q), getattr(self, f"k_norm{suffix}")(k), v
+
+    def forward(self, x, cos, sin, batch: AttnBatch, arena: KVArena):
+        q, k, v = _split(lambda t: self._qkv(t, ""), lambda t: self._qkv(t, "_moe_gen"), x, batch.n_und)
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+        ak, av = arena.k[self.layer_idx], arena.v[self.layer_idx]
+        ak[batch.slots] = k
+        av[batch.slots] = v
+        nc = batch.n_causal
+        if batch.full is None:
+            out = _attend(q, ak, av, batch.causal, (-1, 0), self.scale)
+        elif batch.causal is None:
+            out = _attend(q, ak, av, batch.full, (-1, -1), self.scale)
+        else:
+            out = torch.cat([_attend(q[:nc], ak, av, batch.causal, (-1, 0), self.scale),
+                             _attend(q[nc:], ak, av, batch.full, (-1, -1), self.scale)])
+        out = out.view(-1, self.n_heads * self.head_dim)
+        return _split(self.o_proj, self.o_proj_moe_gen, out, batch.n_und)
+
+
+class MoTLayer(nn.Module):
+    def __init__(self, cfg, layer_idx):
+        super().__init__()
+        h = cfg["hidden_size"]
+        self.self_attn = MoTAttention(h, cfg["num_attention_heads"], cfg["num_key_value_heads"],
+                                      h // cfg["num_attention_heads"], cfg["rms_norm_eps"], layer_idx)
+        self.mlp = SwiGLU(h, cfg["intermediate_size"])
+        self.mlp_moe_gen = SwiGLU(h, cfg["intermediate_size"])
+        for name in ("input_layernorm", "post_attention_layernorm"):
+            setattr(self, name, RMSNorm(h, cfg["rms_norm_eps"]))
+            setattr(self, name + "_moe_gen", RMSNorm(h, cfg["rms_norm_eps"]))
+
+    def forward(self, x, cos, sin, batch, arena):
+        h = _split(self.input_layernorm, self.input_layernorm_moe_gen, x, batch.n_und)
+        x = x + self.self_attn(h, cos, sin, batch, arena)
+        h = _split(lambda t: self.mlp(self.post_attention_layernorm(t)),
+                   lambda t: self.mlp_moe_gen(self.post_attention_layernorm_moe_gen(t)), x, batch.n_und)
+        return x + h
+
+
+class _Decoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(cfg["vocab_size"], cfg["hidden_size"])
+        self.layers = nn.ModuleList([MoTLayer(cfg, i) for i in range(cfg["num_hidden_layers"])])
+        self.norm = RMSNorm(cfg["hidden_size"], cfg["rms_norm_eps"])
+        self.norm_moe_gen = RMSNorm(cfg["hidden_size"], cfg["rms_norm_eps"])
+
+
+class MoTLanguageModel(nn.Module):
+    """Qwen2-architecture MoT decoder (understanding + generation experts)."""
+
+    def __init__(self, cfg: dict):
+        super().__init__()
+        self.cfg = cfg
+        self.model = _Decoder(cfg)
+        self.lm_head = nn.Linear(cfg["hidden_size"], cfg["vocab_size"], bias=False)
+        head_dim = cfg["hidden_size"] // cfg["num_attention_heads"]
+        self.rotary = RotaryEmbedding(head_dim, cfg.get("rope_theta", 1e6))
+
+    @property
+    def num_layers(self):
+        return len(self.model.layers)
+
+    def forward(self, x, positions, batch: AttnBatch, arena: KVArena, cos_sin=None):
+        cos, sin = cos_sin if cos_sin is not None else self.rotary(positions, x.dtype)
+        for layer in self.model.layers:
+            x = layer(x, cos, sin, batch, arena)
+        return _split(self.model.norm, self.model.norm_moe_gen, x, batch.n_und)
+
+
+# --------------------------------------------------------- SigLIP NaViT (BAGEL)
+class FP32LayerNorm(nn.LayerNorm):
+    """LayerNorm computed in float32 (as under autocast), returned in input dtype."""
+
+    def forward(self, x):
+        return F.layer_norm(x.float(), self.normalized_shape, self.weight.float(), self.bias.float(),
+                            self.eps).to(x.dtype)
+
+
+class _SiglipAttention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads, self.head_dim = heads, dim // heads
+        self.q_proj, self.k_proj, self.v_proj = (nn.Linear(dim, dim) for _ in range(3))
+        self.out_proj = nn.Linear(dim, dim)
+
+    def forward(self, x, cu, max_len):
+        T = x.shape[0]
+        q = self.q_proj(x).view(T, self.heads, self.head_dim)
+        k = self.k_proj(x).view(T, self.heads, self.head_dim)
+        v = self.v_proj(x).view(T, self.heads, self.head_dim)
+        out = varlen_attn(q, k, v, cu, cu, max_len, max_len, window_size=(-1, -1))
+        return self.out_proj(out.reshape(T, -1))
+
+
+class _SiglipMLP(nn.Module):
+    def __init__(self, dim, inter):
+        super().__init__()
+        self.fc1, self.fc2 = nn.Linear(dim, inter), nn.Linear(inter, dim)
+
+    def forward(self, x):
+        return self.fc2(F.gelu(self.fc1(x), approximate="tanh"))
+
+
+class _SiglipLayer(nn.Module):
+    def __init__(self, dim, heads, inter, eps):
+        super().__init__()
+        self.layer_norm1, self.layer_norm2 = FP32LayerNorm(dim, eps=eps), FP32LayerNorm(dim, eps=eps)
+        self.self_attn = _SiglipAttention(dim, heads)
+        self.mlp = _SiglipMLP(dim, inter)
+
+    def forward(self, x, cu, max_len):
+        x = x + self.self_attn(self.layer_norm1(x), cu, max_len)
+        return x + self.mlp(self.layer_norm2(x))
+
+
+class _SiglipEmbeddings(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        p = cfg["patch_size"]
+        self.patch_embedding = nn.Linear(cfg.get("num_channels", 3) * p * p, cfg["hidden_size"])
+        self.position_embedding = nn.Embedding((cfg["image_size"] // p) ** 2, cfg["hidden_size"])
+
+
+class _SiglipEncoder(nn.Module):
+    def __init__(self, cfg, n_layers):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            _SiglipLayer(cfg["hidden_size"], cfg["num_attention_heads"], cfg["intermediate_size"],
+                         cfg.get("layer_norm_eps", 1e-6)) for _ in range(n_layers)
+        ])
+
+
+class _SiglipVision(nn.Module):
+    def __init__(self, cfg, n_layers):
+        super().__init__()
+        self.embeddings = _SiglipEmbeddings(cfg)
+        self.encoder = _SiglipEncoder(cfg, n_layers)
+        self.post_layernorm = FP32LayerNorm(cfg["hidden_size"], eps=cfg.get("layer_norm_eps", 1e-6))
+
+
+class SiglipNaViT(nn.Module):
+    """Variable-resolution SigLIP: packed patches of many images per call."""
+
+    def __init__(self, cfg: dict, n_layers: int):
+        super().__init__()
+        self.vision_model = _SiglipVision(cfg, n_layers)
+
+    def forward(self, patches, pos_ids, cu, max_len):
+        vm = self.vision_model
+        x = vm.embeddings.patch_embedding(patches) + vm.embeddings.position_embedding(pos_ids)
+        for layer in vm.encoder.layers:
+            x = layer(x, cu, max_len)
+        return vm.post_layernorm(x)
+
+
+class MLPConnector(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.fc1, self.fc2 = nn.Linear(in_dim, out_dim), nn.Linear(out_dim, out_dim)
+
+    def forward(self, x):
+        return self.fc2(F.gelu(self.fc1(x), approximate="tanh"))
+
+
+def sincos_2d(dim: int, grid: int) -> torch.Tensor:
+    """Fixed 2-D sine-cosine position table of shape [grid*grid, dim]."""
+    def one_d(d, pos):
+        omega = 1.0 / 10000 ** (np.arange(d // 2, dtype=np.float64) / (d / 2.0))
+        out = np.einsum("m,d->md", pos.reshape(-1), omega)
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+    gw, gh = np.meshgrid(np.arange(grid, dtype=np.float32), np.arange(grid, dtype=np.float32))
+    emb = np.concatenate([one_d(dim // 2, gh), one_d(dim // 2, gw)], axis=1)
+    return torch.from_numpy(emb).float()
+
+
+class PositionTable(nn.Module):
+    def __init__(self, grid: int, dim: int):
+        super().__init__()
+        self.pos_embed = nn.Parameter(sincos_2d(dim, grid), requires_grad=False)
+
+    def forward(self, ids):
+        return self.pos_embed[ids]
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden: int, freq_dim: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.Linear(freq_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden))
+        self.freq_dim = freq_dim
+
+    def forward(self, t):
+        half = self.freq_dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+# ----------------------------------------------------------------- FLUX VAE
+class _ResBlock(nn.Module):
+    def __init__(self, cin, cout):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(32, cin, eps=1e-6)
+        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm2 = nn.GroupNorm(32, cout, eps=1e-6)
+        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
+        self.nin_shortcut = nn.Conv2d(cin, cout, 1) if cin != cout else None
+
+    def forward(self, x):
+        h = self.conv2(F.silu(self.norm2(self.conv1(F.silu(self.norm1(x))))))
+        return (x if self.nin_shortcut is None else self.nin_shortcut(x)) + h
+
+
+class _AttnBlock(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.norm = nn.GroupNorm(32, c, eps=1e-6)
+        self.q, self.k, self.v, self.proj_out = (nn.Conv2d(c, c, 1) for _ in range(4))
+
+    def forward(self, x):
+        h = self.norm(x)
+        b, c, hh, ww = h.shape
+        q, k, v = (t(h).reshape(b, 1, c, hh * ww).transpose(-1, -2) for t in (self.q, self.k, self.v))
+        out = F.scaled_dot_product_attention(q, k, v).transpose(-1, -2).reshape(b, c, hh, ww)
+        return x + self.proj_out(out)
+
+
+class _Up(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, 3, padding=1)
+
+    def forward(self, x):
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class _Down(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, 3, stride=2)
+
+    def forward(self, x):
+        return self.conv(F.pad(x, (0, 1, 0, 1)))
+
+
+class _Stage(nn.Module):
+    pass
+
+
+class FluxAutoEncoder(nn.Module):
+    """The FLUX.1 VAE used by BAGEL (16 latent channels, 8x downsampling)."""
+
+    def __init__(self, ch=128, ch_mult=(1, 2, 4, 4), num_res_blocks=2, z_channels=16,
+                 scale_factor=0.3611, shift_factor=0.1159):
+        super().__init__()
+        self.scale_factor, self.shift_factor = scale_factor, shift_factor
+        n = len(ch_mult)
+        # encoder
+        enc = nn.Module()
+        enc.conv_in = nn.Conv2d(3, ch, 3, padding=1)
+        enc.down = nn.ModuleList()
+        in_mult = (1,) + tuple(ch_mult)
+        block_in = ch
+        for i in range(n):
+            stage = _Stage()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_in, block_out = ch * in_mult[i], ch * ch_mult[i]
+            for _ in range(num_res_blocks):
+                stage.block.append(_ResBlock(block_in, block_out))
+                block_in = block_out
+            if i != n - 1:
+                stage.downsample = _Down(block_in)
+            enc.down.append(stage)
+        enc.mid = nn.Module()
+        enc.mid.block_1, enc.mid.attn_1, enc.mid.block_2 = _ResBlock(block_in, block_in), _AttnBlock(block_in), _ResBlock(block_in, block_in)
+        enc.norm_out = nn.GroupNorm(32, block_in, eps=1e-6)
+        enc.conv_out = nn.Conv2d(block_in, 2 * z_channels, 3, padding=1)
+        self.encoder = enc
+        # decoder
+        dec = nn.Module()
+        block_in = ch * ch_mult[-1]
+        dec.conv_in = nn.Conv2d(z_channels, block_in, 3, padding=1)
+        dec.mid = nn.Module()
+        dec.mid.block_1, dec.mid.attn_1, dec.mid.block_2 = _ResBlock(block_in, block_in), _AttnBlock(block_in), _ResBlock(block_in, block_in)
+        ups = [None] * n
+        for i in reversed(range(n)):
+            stage = _Stage()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = ch * ch_mult[i]
+            for _ in range(num_res_blocks + 1):
+                stage.block.append(_ResBlock(block_in, block_out))
+                block_in = block_out
+            if i != 0:
+                stage.upsample = _Up(block_in)
+            ups[i] = stage
+        dec.up = nn.ModuleList(ups)
+        dec.norm_out = nn.GroupNorm(32, block_in, eps=1e-6)
+        dec.conv_out = nn.Conv2d(block_in, 3, 3, padding=1)
+        self.decoder = dec
+        self._n = n
+        self._num_res_blocks = num_res_blocks
+
+    def decode(self, z):
+        d = self.decoder
+        z = z / self.scale_factor + self.shift_factor
+        h = d.conv_in(z)
+        h = d.mid.block_2(d.mid.attn_1(d.mid.block_1(h)))
+        for i in reversed(range(self._n)):
+            for blk in d.up[i].block:
+                h = blk(h)
+            if i != 0:
+                h = d.up[i].upsample(h)
+        return d.conv_out(F.silu(d.norm_out(h)))
+
+    def encode(self, x):
+        e = self.encoder
+        h = e.conv_in(x)
+        for i in range(self._n):
+            for blk in e.down[i].block:
+                h = blk(h)
+            if i != self._n - 1:
+                h = e.down[i].downsample(h)
+        h = e.mid.block_2(e.mid.attn_1(e.mid.block_1(h)))
+        h = e.conv_out(F.silu(e.norm_out(h)))
+        mean = h.chunk(2, dim=1)[0]
+        return self.scale_factor * (mean - self.shift_factor)