MoonViT-SO-400M

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

MoonViT-SO-400M — ONNX

An ONNX export of 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

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)

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)

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):

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)

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 for training details.