---
license: mit
pipeline_tag: image-feature-extraction
library_name: onnxruntime
base_model:
- moonshotai/MoonViT-SO-400M
tags:
- onnx
- vision-encoder
---

<div align="center">
  <img width="30%" src="figures/logo.png">
</div>

# MoonViT-SO-400M — ONNX

An ONNX export of **[moonshotai/MoonViT-SO-400M](https://huggingface.co/moonshotai/MoonViT-SO-400M)**,
the native-resolution vision encoder from Kimi-VL (initialized from and continually pre-trained on
SigLIP-SO-400M). Built with the standalone `MoonVit.py` script (`build` / `eval` / `inference`).

**Parity vs the original PyTorch model: cosine = 1.000000** (see `eval` below).

## What's in this folder

| File | Description |
|---|---|
| `moonvit.onnx` | The exported graph. Legacy export = single file (~1.7 GB); dynamo export = graph + `moonvit.onnx.data` (external weights). |
| `manifest.json` | Export metadata — grid, I/O shapes, normalization, which exporter was used. |
| `README.md` | This file. |

## Important: fixed-grid export

MoonViT packs patches NaViT-style and its internals branch on `grid_hws.tolist()` (per-image
Python loops for the interpolated position embedding, the 2D RoPE, and the 2×2 patch merger), so
the graph is **data-dependent on the image grid**. This export **bakes one fixed grid as a
constant**. The default is **28×28 patches = 392×392 px**; check `manifest.json` for the exact grid
this file was built with. For other input resolutions, export another variant
(`--grid-h/--grid-w`).

- **Input** `pixel_values`: `float32 [L, 3, 14, 14]` — `L = grid_h × grid_w` packed 14×14 patches,
  raster order, normalized with mean/std = 0.5.
- **Output** `image_features`: `float32 [L/4, 4, 1152]` — merged tokens after the 2×2 patch merger
  (e.g. 28×28 → `[196, 4, 1152]`).

## Usage (onnxruntime)

```python
import json, numpy as np, onnxruntime as ort
from PIL import Image

d = "."                                    # this folder
man = json.load(open(f"{d}/manifest.json"))
gh, gw, P = man["grid_h"], man["grid_w"], man["patch_size"]

# preprocess an image to the export grid (raster-order packed patches, mean/std 0.5)
img = Image.open("figures/demo.png").convert("RGB").resize((gw * P, gh * P), Image.BICUBIC)
x = (np.asarray(img, np.float32) / 255.0 - 0.5) / 0.5           # [H,W,3]
x = x.transpose(2, 0, 1).reshape(3, gh, P, gw, P).transpose(1, 3, 0, 2, 4)
pixel_values = np.ascontiguousarray(x.reshape(gh * gw, 3, P, P))

sess = ort.InferenceSession(f"{d}/{man['file']}", providers=["CPUExecutionProvider"])
feats = sess.run(None, {"pixel_values": pixel_values})[0]       # [gh*gw/4, 4, 1152]
print(feats.shape)
```

Or via the script: `uv run MoonVit.py inference --onnx-model-path <this_dir> --image img.png`.

## How it was built (`MoonVit.py`)

```bash
uv run MoonVit.py build     --output MoonVitOnnx                 # legacy exporter (default)
uv run MoonVit.py build     --output MoonVitOnnx --dynamo        # new torch.onnx dynamo exporter
uv run MoonVit.py eval      --onnx-model-path MoonVitOnnx        # original PyTorch vs ONNX
uv run MoonVit.py inference --onnx-model-path MoonVitOnnx --image img.png
```

Both exporters are supported and produce **identical accuracy (cos 1.000000)**:

| Exporter | Graph | Weights | Notes |
|---|---|---|---|
| `legacy` (default) | ~5,050 nodes | inline (~1.7 GB single file) | unrolls the data-dependent loops against the baked grid |
| `dynamo` (`--dynamo`) | ~1,794 nodes | external `.onnx.data` | cleaner/smaller graph; requires `onnxscript` |

Export notes (handled automatically by the script; both verified bit-faithful before export):
1. **Complex RoPE → real cos/sin.** ONNX has no complex dtype; the 2D-RoPE `freqs_cis` is
   precomputed for the baked grid and applied as real-valued rotation.
2. **`PytorchGELUTanh` shim** — the upstream remote code imports it from `transformers.activations`
   (removed in newer transformers); it's exactly `nn.GELU(approximate="tanh")`.
3. **Dynamo only:** full-attention SDPA (single baked image ⇒ the block-diagonal mask is all-True)
   and concrete-int patch-merger / pos-emb interpolation, to avoid unbacked-symint ops.
4. **Legacy only:** int32→int64 normalization of `Slice`/`Gather` index tensors.

## Eval result

```
grid=28x28   original PyTorch vs ONNX
  sample 0: cos=1.000000  max|Δ|=9.995e-02
  sample 1: cos=1.000000  max|Δ|=6.195e-03
  sample 2: cos=1.000000  max|Δ|=2.052e-03
=== PASS (worst cosine 1.000000, tol 0.999) ===
```
(`max|Δ|` is on large-magnitude features, feature std ≈ 4.2 — cosine 1.0 is the real signal.)

## Original model reference (PyTorch)

```python
from PIL import Image
from transformers import AutoModel, AutoImageProcessor

model_path = "moonshotai/MoonViT-SO-400M"
model = AutoModel.from_pretrained(model_path, torch_dtype="auto", device_map="auto",
                                  trust_remote_code=True)
processor = AutoImageProcessor.from_pretrained(model_path, trust_remote_code=True)

image = Image.open("./figures/demo.png")
proc = processor(image, return_tensors="pt").to(dtype=model.dtype, device=model.device)
image_features: list = model(proc.pixel_values, proc.image_grid_hws)
print(image_features[0].dtype, image_features[0].shape)   # e.g. bf16, [N, 4, 1152]
```

See the [Kimi-VL Technical Report](https://huggingface.co/papers/2504.07491) for training details.