# 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()