""" export_tail.py — reproduce siglip2_tail_fp16.onnx from the base model. Exports ONLY the static tail of SigLIP2 NaFlex's vision tower (27 encoder layers + post-layernorm + attention-pooling head), leaving the embeddings in PyTorch. WHY THE SPLIT ------------- Siglip2VisionTransformer.forward is: hidden_states = self.embeddings(pixel_values, spatial_shapes) <-- data-dependent enc_mask = _prepare_4d_attention_mask(attention_mask, dtype) encoder_out = self.encoder(inputs_embeds=hidden_states, attention_mask=enc_mask) last_hidden = self.post_layernorm(encoder_out.last_hidden_state) pooled = self.head(last_hidden, attention_mask) `self.embeddings` interpolates the learned position-embedding grid to EACH image's (h_patches, w_patches), taken from `spatial_shapes`. That is data-dependent control flow: tracing it bakes in whichever aspect ratios were in the export batch, and every other aspect ratio then silently gets the wrong position embeddings. The model keeps running and keeps returning plausible-looking vectors, which is what makes the failure dangerous. Note the shapes are NOT the problem: `pixel_values` is always (B, 576, 768) regardless of aspect ratio. Only the VALUES of `spatial_shapes` vary. So the tail is perfectly static and exports cleanly. Usage: python export_tail.py --out siglip2_tail_fp16.onnx --fp16 python export_tail.py --out siglip2_tail_fp32.onnx # fp32 on CPU python export_tail.py --out t.onnx --fp16 --verify # + held-out check """ import argparse import os import numpy as np import torch import torch.nn as nn from PIL import Image MODEL_ID = "google/siglip2-so400m-patch16-naflex" REVISION = "cc24074f717b612951c2dead130904ab9b65a81e" MAX_NUM_PATCHES = 576 def _expand_mask(mask_2d: torch.Tensor, dtype: torch.dtype, tgt_len: int) -> torch.Tensor: """transformers' _prepare_4d_attention_mask, written with plain ops so it traces.""" b, s = mask_2d.shape m = mask_2d[:, None, None, :].to(dtype).expand(b, 1, tgt_len, s) return (1.0 - m) * torch.finfo(dtype).min class VisionEncoderHead(nn.Module): """The static tail: encoder + post_layernorm + attention-pooling head.""" def __init__(self, vision_model): super().__init__() self.encoder = vision_model.encoder self.post_layernorm = vision_model.post_layernorm self.head = vision_model.head self.num_heads = vision_model.head.num_heads def forward(self, hidden_states, attention_mask): dtype = hidden_states.dtype seq = hidden_states.shape[1] enc_mask = _expand_mask(attention_mask, dtype, seq) eo = self.encoder(inputs_embeds=hidden_states, attention_mask=enc_mask) last_hidden = eo.last_hidden_state if hasattr(eo, "last_hidden_state") else eo[0] last_hidden = self.post_layernorm(last_hidden) # Siglip2MultiheadAttentionPoolingHead, inlined so the mask reshape traces b = last_hidden.shape[0] probe = self.head.probe.repeat(b, 1, 1) tgt = probe.shape[1] m = _expand_mask(attention_mask, dtype, tgt) m = m.repeat(1, self.num_heads, tgt, 1).reshape(-1, tgt, seq) h = self.head.attention(probe, last_hidden, last_hidden, attn_mask=m)[0] h = h + self.head.mlp(self.head.layernorm(h)) return h[:, 0] def probe_images(n, seed=0): """Deterministic MIXED aspect ratios — the export batch should not be uniform.""" rng = np.random.default_rng(seed) sizes = [(93, 209), (327, 425), (466, 480), (304, 480), (223, 480), (189, 480), (320, 480), (240, 400)] return [Image.fromarray(rng.integers(0, 255, (sizes[i % 8][1], sizes[i % 8][0], 3), dtype=np.uint8)) for i in range(n)] def main(): ap = argparse.ArgumentParser() ap.add_argument("--out", default="siglip2_tail_fp16.onnx") ap.add_argument("--batch", type=int, default=8) ap.add_argument("--fp16", action="store_true", help="export fp16 (needs CUDA)") ap.add_argument("--opset", type=int, default=17) ap.add_argument("--verify", action="store_true", help="check the ONNX on aspect ratios NOT in the export batch") args = ap.parse_args() from transformers import AutoModel, AutoProcessor dtype = torch.float16 if args.fp16 else torch.float32 device = "cuda" if args.fp16 else "cpu" if args.fp16 and not torch.cuda.is_available(): raise SystemExit("--fp16 needs CUDA (fp16 tracing on CPU is unsupported for some ops)") print(f"loading {MODEL_ID}@{REVISION[:8]} as {dtype} on {device}") model = AutoModel.from_pretrained(MODEL_ID, revision=REVISION, torch_dtype=dtype).to(device).eval() proc = AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION) imgs = probe_images(args.batch) print(f"export batch aspect ratios: {[im.size for im in imgs]}") inp = proc(images=imgs, return_tensors="pt", max_num_patches=MAX_NUM_PATCHES) inp = {k: v.to(device) for k, v in inp.items()} with torch.no_grad(): hidden = model.vision_model.embeddings(inp["pixel_values"], inp["spatial_shapes"]) mask = inp["pixel_attention_mask"].to(torch.int32) print(f"hidden_states {tuple(hidden.shape)} {hidden.dtype} | " f"attention_mask {tuple(mask.shape)} {mask.dtype}") tail = VisionEncoderHead(model.vision_model).eval() # sanity: the split must reproduce the full model before anything is exported with torch.no_grad(): o = model.get_image_features(**inp) ref = (o if isinstance(o, torch.Tensor) else o.pooler_output).float() got = tail(hidden, inp["pixel_attention_mask"]).float() cos = torch.nn.functional.cosine_similarity(got, ref, dim=-1).min().item() print(f"split vs full model: cosine min {cos:.8f} " f"max abs diff {(got - ref).abs().max().item():.8f}") if cos < 0.9999: raise SystemExit("split does not reproduce the full model — aborting") print(f"exporting -> {args.out} (opset {args.opset})") torch.onnx.export( tail, (hidden, mask), args.out, input_names=["hidden_states", "attention_mask"], output_names=["pooled"], dynamic_axes={"hidden_states": {0: "batch"}, "attention_mask": {0: "batch"}, "pooled": {0: "batch"}}, opset_version=args.opset, do_constant_folding=True) print(f"done: {os.path.getsize(args.out) / 1e6:.0f} MB") if args.verify: # THE check that matters: aspect ratios the tracer never saw. import onnxruntime as ort held = probe_images(4, seed=99) print(f"\nverifying on held-out aspect ratios {[im.size for im in held]}") i2 = proc(images=held, return_tensors="pt", max_num_patches=MAX_NUM_PATCHES) i2 = {k: v.to(device) for k, v in i2.items()} with torch.no_grad(): h2 = model.vision_model.embeddings(i2["pixel_values"], i2["spatial_shapes"]) pt = tail(h2, i2["pixel_attention_mask"]).float().cpu().numpy() sess = ort.InferenceSession(args.out, providers=["CPUExecutionProvider"]) ox = sess.run(["pooled"], { "hidden_states": h2.cpu().numpy(), "attention_mask": i2["pixel_attention_mask"].to(torch.int32).cpu().numpy()})[0] a = ox / np.linalg.norm(ox, axis=1, keepdims=True) b = pt / np.linalg.norm(pt, axis=1, keepdims=True) c = (a * b).sum(1) print(f"cosine(onnx, pytorch) on unseen aspect ratios: min {c.min():.6f}") print("PASS" if c.min() > 0.999 else "FAIL — position embeddings may be baked in") if __name__ == "__main__": main()