embeddinggemma-2-vision-static-onnx

The vision encoder of EmbeddingGemma 2 as ONNX with static shapes, so it compiles and runs on the Intel NPU. The text half is arankays/embeddinggemma-2-text-int8-cw-onnx; its onnx/model_image_<tokens>.onnx graphs take this encoder's output directly.

What changed

onnx-community/embeddinggemma-2-ONNX's vision_encoder.onnx pools patches into soft-token cells, then keeps only the cells holding a real patch with NonZero → Transpose → GatherND. That data-dependent selection is the only dynamic shape, and the NPU compiler refuses it (karanshukla/vinoSearch#2).

These graphs fix the inputs at the processor's padded size and drop those three nodes, so image_features is always every cell. The weights are untouched: onnx/vision_encoder.onnx_data is upstream's fp32 file byte for byte (SHA-256 ca8d28c9…c01e37e). Made by bench/static_vision.py.

All three share onnx/vision_encoder.onnx_data (671 MB).

Usage

1. Preprocess. Run the image through Gemma4ImageProcessor (processor_config.json) with max_soft_tokens set to the graph's token count. The processor already pads to tokens × 9 patches, which is the graph's input shape.
2. Keep the valid rows. Take the patch grid as h, w = max(pixel_position_ids over real patches) + 1. Padding patches have position -1. Keep the first (h // 3) * (w // 3) rows of image_features. The graph's image_mask output is right on CPU but wrong on the NPU, so compute the mask on the host.
3. Embed. Feed the features to the text model with input_ids = [2, 255999, 258880 × N, 258882, 1] (bos, begin-of-image, N image tokens, end-of-image, eos), N being the rows kept. With the text repo's model_image_<tokens>.onnx, you can pass all rows unsliced: it merges the first N and ignores the rest. Pad input_ids and attention_mask at the end with 0 to the graph's length. Use sentence_embedding, L2-normalised. Image and text embeddings share one space, so text queries (task: search result | query: …) search images directly.

import numpy as np
import openvino as ov

core = ov.Core()
vision = core.compile_model("onnx/vision_encoder_static_140.onnx", "NPU", {"CACHE_DIR": "ov-cache"})
# pixel_values [1, 1260, 768], pixel_position_ids [1, 1260, 2] from the processor (max_soft_tokens=140)
feats = vision({"pixel_values": pv, "pixel_position_ids": pos})[vision.outputs[0]]
real = pos[0][(pos[0] >= 0).all(-1)]
h, w = real.max(0) + 1
n = (h // 3) * (w // 3)  # image tokens; feats[:n] are the image's features

Results

Measured on a Core 5 320 (Wildcat Lake NPU), OpenVINO 2026.4.1, NPU driver 1.35.0, on 2026-10-19. Both stages ran on the NPU (this encoder, then the text repo's INT8 model_image_<tokens>.onnx).

The table shows retrieval on Flickr30k, 1k test split (1,000 images, 5,000 captions):

The cosine column compares against sentence-transformers with PyTorch at the same token count. The NPU was 100% busy and the pipeline used no CPU time. Package energy is from RAPL, against a 1.65 W idle.

140 is the suggested default: it is 1.8 R@1 below 280 with the same R@10, at half the time and energy.

NPU notes:

• Use fp32 weights. NNCF INT8_SYM weight compression is no faster on the NPU (the encoder is compute-bound) and breaks it there (cosine 0.67). It works on CPU (0.9996).
• First compile takes time: about 95 s at 280 tokens, 28 s at 140 and 12 s at 70. Use a CACHE_DIR.

Used by vinoSearch.

License

Apache 2.0, same as the base model. EmbeddingGemma 2 is by Google DeepMind. The graphs are modified from onnx-community/embeddinggemma-2-ONNX (inputs fixed, the valid-cell selection removed); the weights are unchanged.