export/eg2_vision.py

# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
#
# Re-authoring of the Gemma 4 vision encoder (Gemma4VisionModel + the EmbeddingGemma 2 vision embedder) as
# described by modeling_gemma4.py / modeling_embedding_gemma2.py in Hugging Face transformers 5.19
# (Copyright 2026 The HuggingFace Team, Apache-2.0). Model weights: Google DeepMind (Gemma Team), Apache-2.0,
# https://huggingface.co/google/embeddinggemma-2 — see ../NOTICE for the attribution and modification list.
"""EmbeddingGemma 2 vision encoder re-authored for Core AI export: patches -> soft tokens in the text space.

One fixed shape (the default budget of 280 soft tokens = 2520 patches of 16x16 px):
  patches    [1, 2520, 768] fp16, integer pixel values 0..255 (exact in fp16), 16x16x3 flattened (py, px, c),
             CELL-MAJOR order (the 9 patches of each 3x3 pooling cell contiguous, cells row-major), zero padded
  pos_x/pos_y [1, 2520] int32 patch coordinates (padding 0)
  cos/sin    [1, 1, 2520, 64] fp16 axial RoPE tables computed on the host (channels 0..31 by x, 32..63 by y)
  patch_mask [1, 2520] fp16, 1 = real patch, 0 = padding
  -> (P - 127.5) -> Linear 768->768 (weights x 2/255) + posx[x] + posy[y] -> 16 encoder layers (12 heads x 64, sandwich RMSNorm, q/k/v
     norms, scale 1.0, bidirectional with a key mask) -> 3x3 average pooling = reshape [280, 9, 768].mean
     -> RMSNorm without scale -> Linear 768->512
  soft_tokens [1, 280, 512] fp16; rows at index >= n_soft (= cells of the image) are padding and must be ignored.

Differences from the upstream graph (all mathematically identical in exact arithmetic):
  - upstream feeds p = P/255 and computes 2p - 1. P/255 is not representable in fp16 (rounding it alone costs a
    heavily upscaled tiny image cosine 0.996), so the graph takes the integer pixel P and computes P - 127.5
    (exact in fp16) and the 2/255 scale is folded into input_proj's weights: W(2P/255 - 1) == (2W/255)(P - 127.5).
  - cell-major patch order makes the 3x3 pooling a reshape + mean instead of a one-hot matmul; attention is
    permutation-equivariant so the layers do not care.
  - the upstream pooler multiplies by sqrt(768) in fp32 before the projector's RMSNorm (and warns that this
    exceeds fp16). RMSNorm(x*s, eps) == RMSNorm(x, eps/s^2), so the scale is dropped and eps/768 used instead.
  - the position-embedding tables are cut to the 840 rows an image of this budget can address (x and y < 840).
  - fp16 "ane" precision mode as in eg2_text.py: every RMSNorm divides the row by its max-abs value before
    squaring (the vision residual reaches ~1400), no divides, finite -1e4 key mask.
"""
from __future__ import annotations

import json
import math
from pathlib import Path

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

PRECISION_MODE = "ane"  # "ane" (all fp16, Neural Engine) | "fp32" (correctness gate of the re-authoring)
MASK_VALUE = -10000.0
NORM_FLOOR = 1e-3
MAX_SOFT = 280
POOL = 3
MAX_PATCHES = MAX_SOFT * POOL * POOL  # 2520
POS_ROWS = MAX_SOFT * POOL  # 840 (budget 280); budget 1120 would need 3360
# Query chunking of the 2520x2520 attention (QUERY_CHUNKS > 1 splits the score matrix into row blocks so that
# no single intermediate exceeds 2520*2520/QUERY_CHUNKS*12 fp16 values); 1 = one matmul per layer.
QUERY_CHUNKS = int(__import__("os").environ.get("EG2_VISION_QCHUNKS", "1"))


class Cfg:
    def __init__(self, model_dir: str):
        c = json.loads((Path(model_dir) / "config.json").read_text())["vision_config"]
        t = json.loads((Path(model_dir) / "config.json").read_text())["text_config"]
        self.hidden = c["hidden_size"]  # 768
        self.inter = c["intermediate_size"]  # 3072
        self.n_layers = c["num_hidden_layers"]  # 16
        self.n_heads = c["num_attention_heads"]  # 12
        self.head_dim = c["head_dim"]  # 64
        self.patch = c["patch_size"]  # 16
        self.pool = c["pooling_kernel_size"]  # 3
        self.eps = c["rms_norm_eps"]  # 1e-6
        self.rope_theta = c["rope_parameters"]["rope_theta"]  # 100
        self.pos_size = c["position_embedding_size"]  # 10240
        self.text_hidden = t["hidden_size"]  # 512
        assert c["hidden_activation"] == "gelu_pytorch_tanh"
        assert c["num_key_value_heads"] == self.n_heads and not c["use_clipped_linears"] and not c["standardize"]
        assert c["rope_parameters"]["rope_type"] == "axial"


def _f16(x: float) -> torch.Tensor:
    return torch.tensor(x, dtype=torch.float16)


def _recip(m: torch.Tensor) -> torch.Tensor:
    r = torch.rsqrt(m)
    return r * r


class RMSNorm(nn.Module):
    """Gemma RMSNorm (weight used as-is), with the max-abs rescaling in "ane" mode."""

    def __init__(self, dim: int, eps: float, with_scale: bool = True):
        super().__init__()
        self.with_scale = with_scale
        if with_scale:
            self.weight = nn.Parameter(torch.ones(dim), requires_grad=False)
        self.eps = eps
        self.register_buffer("sqrt_eps16", _f16(math.sqrt(eps)), persistent=False)
        self.register_buffer("floor16", _f16(NORM_FLOOR), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if PRECISION_MODE == "ane":
            inv_m = _recip(torch.maximum(x.abs().amax(-1, keepdim=True), self.floor16))
            u = x * inv_m
            r = self.sqrt_eps16 * inv_m
            y = u * torch.rsqrt((u * u).mean(-1, keepdim=True) + r * r)
            return y * self.weight if self.with_scale else y
        xf = x.float()
        y = xf * torch.pow(xf.pow(2).mean(-1, keepdim=True) + self.eps, -0.5)
        if self.with_scale:
            y = y * self.weight.float()
        return y.to(x.dtype)


def rope_axial(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x [1, H, N, 64]: rotate_half inside each 32-channel half (x-half, y-half). cos/sin [1, 1, N, 64]."""
    B, H, N, D = x.shape
    xr = x.reshape(B, H, N, 2, D // 2)
    h = D // 4
    rot = torch.cat((-xr[..., h:], xr[..., :h]), dim=-1).reshape(B, H, N, D)
    return x * cos + rot * sin


class Attention(nn.Module):
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.n_heads, self.hd = cfg.n_heads, cfg.head_dim
        self.q_proj = nn.Linear(cfg.hidden, cfg.n_heads * cfg.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.hidden, cfg.n_heads * cfg.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.hidden, cfg.n_heads * cfg.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.hidden, bias=False)
        self.q_norm = RMSNorm(cfg.head_dim, cfg.eps)
        self.k_norm = RMSNorm(cfg.head_dim, cfg.eps)
        self.v_norm = RMSNorm(cfg.head_dim, cfg.eps, with_scale=False)

    def forward(self, x, cos, sin, bias):
        B, N, _ = x.shape
        q = self.q_norm(self.q_proj(x).view(B, N, self.n_heads, self.hd)).transpose(1, 2)
        k = self.k_norm(self.k_proj(x).view(B, N, self.n_heads, self.hd)).transpose(1, 2)
        v = self.v_norm(self.v_proj(x).view(B, N, self.n_heads, self.hd)).transpose(1, 2)
        q, k = rope_axial(q, cos, sin), rope_axial(k, cos, sin)
        kt = k.transpose(-1, -2)
        if QUERY_CHUNKS > 1:
            outs = []
            for qc in q.chunk(QUERY_CHUNKS, dim=2):
                s = torch.matmul(qc, kt) + bias  # scale 1.0
                p = torch.softmax(s, dim=-1) if PRECISION_MODE == "ane" else torch.softmax(s.float(), dim=-1).to(q.dtype)
                outs.append(torch.matmul(p, v))
            o = torch.cat(outs, dim=2)
        else:
            s = torch.matmul(q, kt) + bias  # [1, 12, N, N], scale 1.0
            p = torch.softmax(s, dim=-1) if PRECISION_MODE == "ane" else torch.softmax(s.float(), dim=-1).to(q.dtype)
            o = torch.matmul(p, v)
        return self.o_proj(o.transpose(1, 2).reshape(B, N, self.n_heads * self.hd))


class MLP(nn.Module):
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.hidden, cfg.inter, bias=False)
        self.up_proj = nn.Linear(cfg.hidden, cfg.inter, bias=False)
        self.down_proj = nn.Linear(cfg.inter, cfg.hidden, bias=False)

    def forward(self, x):
        return self.down_proj(F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x))


class Layer(nn.Module):
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.self_attn = Attention(cfg)
        self.mlp = MLP(cfg)
        self.input_layernorm = RMSNorm(cfg.hidden, cfg.eps)
        self.post_attention_layernorm = RMSNorm(cfg.hidden, cfg.eps)
        self.pre_feedforward_layernorm = RMSNorm(cfg.hidden, cfg.eps)
        self.post_feedforward_layernorm = RMSNorm(cfg.hidden, cfg.eps)

    def forward(self, x, cos, sin, bias):
        x = x + self.post_attention_layernorm(self.self_attn(self.input_layernorm(x), cos, sin, bias))
        return x + self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(x)))


class EmbeddingGemma2Vision(nn.Module):
    def __init__(self, cfg: Cfg):
        super().__init__()
        self.cfg = cfg
        d = cfg.patch * cfg.patch * 3
        self.input_proj = nn.Linear(d, cfg.hidden, bias=False)
        self.pos_x = nn.Embedding(POS_ROWS, cfg.hidden)
        self.pos_y = nn.Embedding(POS_ROWS, cfg.hidden)
        self.layers = nn.ModuleList([Layer(cfg) for _ in range(cfg.n_layers)])
        # projector: RMSNorm without scale with eps/768 (replaces x*sqrt(768) followed by RMSNorm(eps)), Linear
        self.pre_projection_norm = RMSNorm(cfg.hidden, cfg.eps / cfg.hidden, with_scale=False)
        self.embedding_projection = nn.Linear(cfg.hidden, cfg.text_hidden, bias=False)
        self.register_buffer("mask_value", _f16(MASK_VALUE), persistent=False)
        self.register_buffer("one", _f16(1.0), persistent=False)
        self.register_buffer("c_mid", _f16(127.5), persistent=False)
        self.register_buffer("inv_pool", _f16(1.0 / (POOL * POOL)), persistent=False)

    def forward(self, patches, pos_x, pos_y, cos, sin, patch_mask):
        B, N, _ = patches.shape
        # upstream: input_proj(2 * P/255 - 1); here P - 127.5 (exact in fp16) with 2/255 folded into input_proj
        x = self.input_proj(patches - self.c_mid) + self.pos_x(pos_x) + self.pos_y(pos_y)
        bias = ((self.one - patch_mask) * self.mask_value)[:, None, None, :]  # [1, 1, 1, N] key mask
        for layer in self.layers:
            x = layer(x, cos, sin, bias)
        pooled = x.view(B, N // (POOL * POOL), POOL * POOL, self.cfg.hidden)
        if PRECISION_MODE == "ane":
            pooled = pooled.sum(2) * self.inv_pool  # 9 values of <= ~1400 each: sum < 13k, fp16-safe
        else:
            pooled = pooled.float().mean(2).to(x.dtype)
        return self.embedding_projection(self.pre_projection_norm(pooled))  # [1, 280, 512]


def load_model(model_dir: str, dtype: torch.dtype) -> EmbeddingGemma2Vision:
    from safetensors import safe_open

    cfg = Cfg(model_dir)
    m = EmbeddingGemma2Vision(cfg)
    sd = {}
    with safe_open(str(Path(model_dir) / "model.safetensors"), "pt") as f:
        for k in f.keys():
            if k.startswith("vision_tower.encoder.layers."):
                nk = k[len("vision_tower.encoder."):].replace(".linear.weight", ".weight")
                sd[nk] = f.get_tensor(k).float()
            elif k == "vision_tower.patch_embedder.input_proj.weight":
                sd["input_proj.weight"] = f.get_tensor(k).float() * (2.0 / 255.0)  # takes P - 127.5, P in 0..255
            elif k == "vision_tower.patch_embedder.position_embedding_table":
                t = f.get_tensor(k).float()
                sd["pos_x.weight"] = t[0, :POS_ROWS].clone()
                sd["pos_y.weight"] = t[1, :POS_ROWS].clone()
            elif k == "embed_vision.embedding_projection.weight":
                sd["embedding_projection.weight"] = f.get_tensor(k).float()
    sd = {k: v.to(dtype) for k, v in sd.items()}
    missing, unexpected = m.load_state_dict(sd, strict=False)
    missing = [k for k in missing if not any(s in k for s in ("sqrt_eps16", "floor16", "mask_value", "one", "c_mid", "inv_pool"))]
    assert not missing and not unexpected, (missing, unexpected)
    return m.to(dtype).eval()