wesjos/embeddinggemma-2-onnx-fix

EmbeddingGemma 2 — ONNX

google/embeddinggemma-2 exported to ONNX, three towers, runs on onnxruntime.

Tower	File	Size	Output
Text	onnx/model.onnx	1.08 GB	sentence_embedding [b, 768]
Vision	onnx/vision_encoder.onnx	671 MB	image_features [n, 512]
Audio	onnx/audio_encoder.onnx	1.17 GB	audio_features [n, 512]

Usage

pip install onnxruntime tokenizers numpy pillow
python demo.py

dog.jpeg x captions:
  a dog                            0.725
  a panda                          0.576
  a penguin using ubuntu computer  0.506

• Outputs are L2-normalized, so a dot product is the cosine similarity.
• Retrieval prefixes: queries take task: search result | query: , documents take title: none | text: .
• Images: vision tower produces [N, 512] soft tokens, which the text tower fuses. Put <|image|> (258880) placeholders in input_ids and pass the features via image_features. Audio works the same way with <|audio|> (258881); video frames each go through the vision tower.
• Image preprocessing (aspect-ratio resize to a multiple of 3*16, 16x16 patchify, pad to 280*3²=2520 patches with pixel_position_ids = -1) is inlined in demo.py and only needs numpy and PIL.

Differences from the upstream export

Upstream onnx-community/embeddinggemma-2-ONNX keeps 48 com.microsoft.RotaryEmbedding nodes in the text tower and 16 com.microsoft.MultiHeadAttention nodes in the vision tower. Those contrib ops only have a correct definition in the ORT kernel; other engines (OpenVINO included) silently mistranslate them into a degenerate rotary and a shifted MHA. The output still looks like a vector, it is just the wrong one.

Here both are replaced with standard-domain explicit ops, so the ONNX and OpenVINO repos agree and can be cross-checked. See tools/fix_rotary.py and tools/fix_mha.py.

The multimodal fusion is baked into the text tower: <|image|> (258880) and <|audio|> (258881) placeholders in input_ids are replaced positionally by the encoder features.

Notes

• The processor classes (Gemma4ImageProcessor, EmbeddingGemma2Processor) need transformers >= 5.18. Inference does not: demo.py reads tokenizer.json directly.
• fp32 weights, so roughly twice the size of the fp16 IR with no precision loss.

Unofficial conversion; see this repo's acceptance numbers. Original model by Google (Apache-2.0).