Upload convert_siglip2.py with huggingface_hub

convert_siglip2.py

@@ -0,0 +1,254 @@
+"""Convert timm SigLIP2 (ViT-B/16, 224) image tower to a GPU-clean, GPU-correct
+LiteRT .tflite for the ML Drift GPU delegate.
+
+SigLIP2 (Google 2025, Apache-2.0) is a SOTA CLIP-style image tower. timm exposes
+it as `vit_base_patch16_siglip_224.v2_webli` (93M params, conv patch-embed, NO
+rope, NO cls token, attention-pool head). The text tower for zero-shot is
+open_clip `ViT-B-16-SigLIP2` (same 768-d space, prompt "a photo of a {label}").
+
+Re-authoring (all verbatim, weights copied, corr ~1.0 vs PyTorch) -- the same set
+proven on PE-Core, minus the rope step (SigLIP2 has no rope):
+  * Attention (x12): fused qkv -> 5D head-split (the GPU "C12" wall). Decompose to
+    separate q/k/v, hand the 4D q/k/v to scaled_dot_product_attention (its lowering
+    keeps the batch-matmul 3D with a materialized transpose -> GPU-resident).
+  * AttentionPoolLatent: single constant-latent query -> a batch-matmul there is
+    const@non-const (rejected / mis-computed). Express as broadcast-multiply +
+    reduce-sum (exact for latent_len=1, GPU-correct).
+  * LayerNorm -> overflow-safe LayerNorm: the delegate reduces the variance in fp16
+    even for an fp32 graph; deep-ViT massive activations overflow fp16 (sum > 65504)
+    -> wrong norm that compounds with depth while still reporting full residency.
+    Scale-before-square keeps the sum in range.
+
+I/O: input [1,3,224,224] NCHW float32 normalized to [-1,1] ((x/255-0.5)/0.5),
+output [1,768] L2-normalized image embedding.
+
+    ~/clipconv/bin/python scripts/convert_siglip2.py
+"""
+import os
+import sys
+import types
+import collections
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import _stub  # noqa: F401  (macOS scipy/_propack guard, import FIRST)
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+
+MODEL = "vit_base_patch16_siglip_224.v2_webli"
+IMG = 224
+OUT_DIR = os.path.expanduser("~/code/litertlm-convert/out/siglip2")
+os.makedirs(OUT_DIR, exist_ok=True)
+FP32 = os.path.join(OUT_DIR, "siglip2_base_224.tflite")
+FP16 = os.path.join(OUT_DIR, "siglip2_base_224_fp16.tflite")
+
+BANNED = {"GATHER_ND", "GATHER", "TOPK_V2", "FLEX_ERF", "ERF", "BROADCAST_TO"}
+
+
+# -------------------------------------------------- overflow-safe LayerNorm
+class SafeLayerNorm(nn.Module):
+    """LayerNorm whose variance reduction can't overflow fp16 (see module docstring)."""
+    SC = 0.03125  # 1/32
+
+    def __init__(self, ln: nn.LayerNorm):
+        super().__init__()
+        self.weight, self.bias, self.eps = ln.weight, ln.bias, ln.eps
+
+    def forward(self, x):
+        xc = x - x.mean(-1, keepdim=True)
+        xs = xc * self.SC
+        var = (xs * xs).mean(-1, keepdim=True) / (self.SC * self.SC)
+        return xc * torch.rsqrt(var + self.eps) * self.weight + self.bias
+
+
+def patch_layernorm(module):
+    for name, child in module.named_children():
+        if isinstance(child, nn.LayerNorm):
+            setattr(module, name, SafeLayerNorm(child))
+        else:
+            patch_layernorm(child)
+
+
+# ------------------------------------------ block Attention -> 4D (no rope)
+def _attn_forward(self, x, *args, **kwargs):
+    B, N, C = x.shape
+    H, d = self.num_heads, self.head_dim
+    q = self.q_proj_d(x).reshape(B, N, H, d).transpose(1, 2)
+    k = self.k_proj_d(x).reshape(B, N, H, d).transpose(1, 2)
+    v = self.v_proj_d(x).reshape(B, N, H, d).transpose(1, 2)
+    q, k = self.q_norm(q), self.k_norm(k)  # Identity for SigLIP2
+    out = F.scaled_dot_product_attention(q, k, v)
+    out = out.transpose(1, 2).reshape(B, N, H * d)
+    out = self.norm(out)  # Identity (scale_norm off)
+    return self.proj(out)
+
+
+def reauthor_attn(attn):
+    C = attn.qkv.in_features
+    w = attn.qkv.weight.data
+    b = attn.qkv.bias.data if attn.qkv.bias is not None else None
+    has_b = b is not None
+    q_proj = nn.Linear(C, C, bias=has_b)
+    k_proj = nn.Linear(C, C, bias=has_b)
+    v_proj = nn.Linear(C, C, bias=has_b)
+    with torch.no_grad():
+        q_proj.weight.copy_(w[:C])
+        k_proj.weight.copy_(w[C:2 * C])
+        v_proj.weight.copy_(w[2 * C:])
+        if has_b:
+            q_proj.bias.copy_(b[:C])
+            k_proj.bias.copy_(b[C:2 * C])
+            v_proj.bias.copy_(b[2 * C:])
+    attn.q_proj_d, attn.k_proj_d, attn.v_proj_d = q_proj, k_proj, v_proj
+    attn.forward = types.MethodType(_attn_forward, attn)
+
+
+# ------------------------------------------ AttentionPoolLatent -> broadcast-reduce
+def _attn_pool_forward(self, x, attn_mask=None):
+    B, N, C = x.shape
+    H, d, L = self.num_heads, self.head_dim, self.latent_len
+    k = self.k_norm(self.k_proj_d(x).reshape(B, N, H, d).transpose(1, 2))  # [B,H,N,d]
+    v = self.v_proj_d(x).reshape(B, N, H, d).transpose(1, 2)               # [B,H,N,d]
+    qc = self.q_const  # [H, L, d] constant, q_norm'd + scaled
+    scores = (qc.unsqueeze(0) * k).sum(dim=-1)        # [B, H, N]
+    attn = scores.softmax(dim=-1).unsqueeze(-1)       # [B, H, N, 1]
+    out = (attn * v).sum(dim=2).reshape(B, L, C)      # [B, L, C]
+    out = self.proj(out)
+    if self.mlp is not None:
+        out = out + self.mlp(self.norm(out))
+    if self.pool == "token":
+        out = out[:, 0]
+    elif self.pool == "avg":
+        out = out.mean(1)
+    return out
+
+
+def reauthor_attn_pool(ap):
+    assert ap.pos_embed is None, "attn_pool pos_embed not handled"
+    C = ap.kv.in_features
+    inner = ap.num_heads * ap.head_dim
+    has_b = ap.kv.bias is not None
+    k_proj = nn.Linear(C, inner, bias=has_b)
+    v_proj = nn.Linear(C, inner, bias=has_b)
+    with torch.no_grad():
+        k_proj.weight.copy_(ap.kv.weight.data[:inner])
+        v_proj.weight.copy_(ap.kv.weight.data[inner:])
+        if has_b:
+            k_proj.bias.copy_(ap.kv.bias.data[:inner])
+            v_proj.bias.copy_(ap.kv.bias.data[inner:])
+        H, d, L = ap.num_heads, ap.head_dim, ap.latent_len
+        ql = ap.q(ap.latent.expand(1, -1, -1)).reshape(1, L, H, d).transpose(1, 2)
+        ql = ap.q_norm(ql) * ap.scale
+    ap.k_proj_d, ap.v_proj_d = k_proj, v_proj
+    ap.register_buffer("q_const", ql.reshape(H, L, d).detach())
+    ap.forward = types.MethodType(_attn_pool_forward, ap)
+
+
+# ------------------------------------------------------------------- wrapper
+class SigLIP2ImageEncoder(nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.m = m
+
+    def forward(self, pixel):
+        m = self.m
+        x = m.patch_embed(pixel)
+        if x.dim() == 4:
+            x = x.flatten(1, 2)
+        if m.pos_embed is not None:
+            x = x + m.pos_embed       # SigLIP2 has no cls token
+        x = m.norm_pre(x)
+        for blk in m.blocks:
+            x = blk(x)
+        x = m.norm(x)
+        x = m.attn_pool(x)            # -> [B, 768]
+        return F.normalize(x, dim=-1)
+
+
+def op_hist(path):
+    from ai_edge_litert.interpreter import Interpreter
+    it = Interpreter(model_path=path)
+    it.allocate_tensors()
+    hist = collections.Counter(d["op_name"] for d in it._get_ops_details())
+    over4d = sum(1 for d in it.get_tensor_details() if len(d.get("shape", [])) > 4)
+    return hist, over4d, it
+
+
+def tflite_run(it, x_nchw):
+    inp = it.get_input_details()[0]
+    shp = list(inp["shape"])
+    x = x_nchw if shp[1] == 3 else np.transpose(x_nchw, (0, 2, 3, 1)).copy()
+    it.set_tensor(inp["index"], x.astype(inp["dtype"]))
+    it.invoke()
+    return it.get_tensor(it.get_output_details()[0]["index"]).astype("float64").reshape(-1)
+
+
+def main():
+    torch.manual_seed(0)
+    print(f"loading {MODEL} (pretrained, apache-2.0) ...")
+    m = timm.create_model(MODEL, pretrained=True, num_classes=0).eval()
+
+    x = torch.randn(1, 3, IMG, IMG)
+    with torch.no_grad():
+        ref = F.normalize(m(x), dim=-1).numpy().flatten()  # original trunk embedding
+
+    for blk in m.blocks:
+        reauthor_attn(blk.attn)
+    reauthor_attn_pool(m.attn_pool)
+    patch_layernorm(m)
+    enc = SigLIP2ImageEncoder(m).eval()
+
+    with torch.no_grad():
+        got = enc(x).numpy().flatten()
+    corr = float(np.corrcoef(ref, got)[0, 1])
+    print(f"EAGER parity (orig vs re-authored): corr {corr:.8f}  max|diff| {np.abs(ref-got).max():.3e}")
+    assert corr > 0.9999, "re-authoring changed the math -- fix before convert"
+
+    print("converting (litert_torch) ...")
+    import litert_torch
+    litert_torch.convert(enc, (x,)).export(FP32)
+
+    hist, over4d, it = op_hist(FP32)
+    bad = {k: v for k, v in hist.items() if k in BANNED}
+    print(f"FP32 ops: {dict(sorted(hist.items(), key=lambda kv: -kv[1]))}")
+    print(f"banned: {bad or 'NONE'} | >4D tensors: {over4d}")
+    o = tflite_run(it, x.numpy())
+    print(f"PARITY tflite(fp32) vs torch: corr {np.corrcoef(ref, o)[0,1]:.6f}")
+    assert not bad and over4d == 0, "GPU blockers remain -- inspect op histogram"
+
+    print("quantizing fp16 (FLOAT_CASTING) ...")
+    from ai_edge_quantizer import quantizer, recipe_manager
+    from ai_edge_quantizer.recipe import AlgorithmName, qtyping
+    rm = recipe_manager.RecipeManager()
+    rm.add_quantization_config(
+        regex=".*",
+        operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
+        op_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=qtyping.TensorQuantizationConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT),
+            compute_precision=qtyping.ComputePrecision.FLOAT,
+        ),
+        algorithm_key=AlgorithmName.FLOAT_CASTING,
+    )
+    if os.path.exists(FP16):
+        os.remove(FP16)
+    qt = quantizer.Quantizer(float_model=FP32)
+    qt.load_quantization_recipe(rm.get_quantization_recipe())
+    qt.quantize().export_model(FP16)
+
+    s32, s16 = os.path.getsize(FP32) / 1e6, os.path.getsize(FP16) / 1e6
+    print(f"SIZE fp32 {s32:.1f} MB -> fp16 {s16:.1f} MB ({s16/s32*100:.0f}%)")
+    h16, o16d, it16 = op_hist(FP16)
+    bad16 = {k: v for k, v in h16.items() if k in BANNED}
+    print(f"FP16 banned: {bad16 or 'NONE'} | >4D: {o16d}")
+    o16 = tflite_run(it16, x.numpy())
+    print(f"PARITY tflite(fp16) vs torch: corr {np.corrcoef(ref, o16)[0,1]:.6f}  "
+          f"fp16-vs-fp32 corr {np.corrcoef(o, o16)[0,1]:.6f}")
+    print("\nDONE:", FP16)
+
+
+if __name__ == "__main__":
+    main()