upshift/siglip2-so400m-patch16-naflex-onnx

siglip2-so400m-patch16-naflex — ONNX (browser-ready)

ONNX export of google/siglip2-so400m-patch16-naflex (revision cc24074), split into a vision and a text model, for onnxruntime-web on WebGPU and wasm (and any other onnxruntime backend). fp32 and fp16, opset 18.

File	Precision	Size	Inputs → output
vision_model.onnx + vision_model.onnx_data	fp32	1.71 GB	pixel_values float32 (B, N, 768), pixel_attention_mask int64 (B, N), spatial_shapes int64 (B, 2) → image_embeds float32 (B, 1152)
text_model.onnx + text_model.onnx_data	fp32	2.83 GB	input_ids int64 (B, 64) → text_embeds float32 (B, 1152)
vision_model_fp16.onnx + vision_model_fp16.onnx_data	fp16	889 MB	same as fp32
text_model_fp16.onnx + text_model_fp16.onnx_data	fp16	1.42 GB	same as fp32

B (batch) and N (number of patches, up to 256) are dynamic. Embeddings are not normalised. config.json, preprocessor_config.json and the tokenizer files are copied from the original repo; export_meta.json holds logit_scale (already exponentiated) and logit_bias for scoring.

fp16 variants (*_fp16.onnx) are for onnxruntime-web WebGPU (needs the shader-f16 feature): half the download, faster inference. Inputs and outputs stay float32, so they are drop-in replacements. The vision model keeps its embeddings and attention-pooling head in fp32. WebGPU computes in half precision, so expect small score shifts: in our tests image embeddings matched fp32 with cosine ≥ 0.9995 and label probabilities moved by up to ~2 points. On wasm/CPU use the fp32 files.

What differs from a plain export

The NaFlex vision tower resizes its position-embedding grid to each image's patch grid with F.interpolate(..., antialias=True) inside a Python loop over the batch. Exported naively, that bakes in one image shape and produces an antialiased Resize that onnxruntime-web's WebGPU backend does not support. In this export the resize is rewritten as separable antialiased-bilinear weight matrices (Wy @ grid @ Wxᵀ) computed in the graph from spatial_shapes: no Loop, no Resize, any aspect ratio.

Parity

fp32

Compared with the PyTorch model on onnxruntime (CPUExecutionProvider), threshold cosine ≥ 0.9999: passed.

Tower	Inputs	Min cosine	Max abs diff
vision	4	1.0000000	3.5e-05
text	6	1.0000000	4.2e-05

Max logit difference: 1.1e-04.

Versions: torch 2.14.1, transformers 5.19.0, onnxruntime 1.30.0.

fp16

Compared with the PyTorch model on onnxruntime (CPUExecutionProvider), threshold cosine ≥ 0.999: passed.

Tower	Inputs	Min cosine	Max abs diff
vision	4	0.9999029	2.3e-02
text	6	0.9999981	9.8e-03

Max logit difference: 4.3e-02.

Versions: torch 2.14.1, transformers 5.19.0, onnxruntime 1.30.0.

Inputs

Image (NaFlex). Resize the image (bilinear, antialiased when downscaling) to a height and width that are multiples of 16, keeping the aspect ratio as closely as possible, with (h/16)·(w/16) ≤ 256 (see get_image_size_for_max_num_patches in transformers). Rescale to [-1, 1] (x / 127.5 - 1), cut into 16×16 patches in row-major order, flatten each patch as (row, column, channel) into 768 values, and pad to N patches. pixel_attention_mask is 1 for real patches, spatial_shapes is (rows, cols) of the patch grid. Use Siglip2ImageProcessor from transformers as the reference.

Text. Lowercase the text, tokenise with the included Gemma tokenizer.json, append <eos> and pad with <pad> (id 0) to exactly 64 tokens. The model pools the last token, so the padding is required. In Python, use Siglip2Tokenizer (transformers 5.x AutoTokenizer returns a GemmaTokenizer that does not lowercase).

Scoring. p = sigmoid(logit_scale · cos(image, text) + logit_bias), per label. SigLIP scores labels independently, so probabilities do not sum to 1; a correct label at 10–40 % with the rest near 0 is normal.

Usage with onnxruntime-web

import * as ort from 'onnxruntime-web/webgpu';

const base = 'https://huggingface.co/upshift/siglip2-so400m-patch16-naflex-onnx/resolve/main/';
const vision = await ort.InferenceSession.create(base + 'vision_model.onnx', {
  executionProviders: ['webgpu', 'wasm'],
  externalData: [{ path: 'vision_model.onnx_data', data: base + 'vision_model.onnx_data' }],
});
const { image_embeds } = await vision.run({
  pixel_values: new ort.Tensor('float32', pixelValues, [1, numPatches, 768]),
  pixel_attention_mask: new ort.Tensor('int64', pixelAttentionMask, [1, numPatches]),
  spatial_shapes: new ort.Tensor('int64', BigInt64Array.from([rows, cols].map(BigInt)), [1, 2]),
});

The text model loads the same way with text_model.onnx / text_model.onnx_data and an input_ids tensor of shape [B, 64].

Limitations

• The text model (1.42 GB in fp16) has to be downloaded and held in memory. If you only need image embeddings, load just the vision model.
• text_model.onnx_data (2.83 GB) is larger than browsers load reliably (ArrayBuffer and wasm memory limits of roughly 2–4 GB). In the browser use the smaller fp16 files on WebGPU; the large files are for onnxruntime in Node, Python or native code.

License

Apache-2.0, as the original model. Original work by Google; see the SigLIP 2 paper and the original model card.