Xiaomi-OCR-0 · Apple Neural Engine edition

Xiaomi-OCR-0 (0.8B OCR VLM) exported to Core ML so that the vision tower, the prefill pass and the token loop all execute on Apple silicon. Nothing was fine-tuned; greedy tokens are identical to the PyTorch reference.

Why this exists (and what it is not)

This is not the fastest way to run this model — a Metal/GPU stack is several times faster (numbers below). It exists for three properties a GPU stack cannot give you:

Zero third-party runtime. Only Apple's Core ML framework — no torch, no ONNX Runtime GPU stack, no llama.cpp/MLX binary to ship, sign and keep updated. Core ML ships with the OS.
The compute does not compete for CPU or GPU. The 24-layer text stack (prefill + decode) runs on the Neural Engine, and the vision tower on the GPU only when the page is big enough to be worth it — so the CPU stays free for your app and the GPU for rendering / other inference. Measured whole-loop host overhead of the reference Python runner: ≈0.65 core (numpy state copies + a 248k-vocab argmax). That is host bookkeeping, not model compute; a native Swift integration holding the 48 state tensors in MLState removes nearly all of it.
Silent, offline, batchable. No network, no daemon, no IPC, no display: load the .mlpackage set, feed an image, get text. Fits a menu-bar agent, a background OCR queue, a folder watcher that transcribes scans overnight, or an iOS share-sheet extension.

Power draw is low by architecture (ANE is the low-power engine, CPU/GPU stay idle); watts were not measured — treat that as a design property, not a benchmark number.

Use a GPU stack (MLX / llama.cpp) when latency is the requirement; use this when footprint, power and zero dependencies are.

Measured (M1 Pro, macOS 26)

Receipt page, 5916 patches at native resolution → 2304-patch bucket:

stage	engine	time
vision	GPU (bucketed 64×36)	0.68 s
prefill (597 tokens)	ANE	1.99 s
decode	ANE	20.6–25.9 tok/s (KV 4096 / 2048)
end-to-end		≈ 4.6–5.5 s / page

Same weights, measured the same way, for honesty:

stack	decode	note
MLX (Metal GPU), community 4-bit	406 tok/s (receipt), 183 tok/s (dense page)	fastest; needs the MLX stack
this package (Core ML)	20–26 tok/s	no dependencies, low power
ONNX int4 on CPU	15.7 tok/s	CPU-bound

Quality: qianziwen is byte-identical to the CPU/ONNX path; receipt differs by one space and one comma; the dense formula page matches the official reference transcript.

Vision buckets

The ANE needs static shapes, so the tower ships as five fixed grids (gh×gw, both even because spatial_merge_size = 2, patches ≤ 2304 — the ANE refuses 3136 with status=0x16):

model	patches	ANE	GPU	ONNX-CPU
`vis_32x32`	1024	304 ms	297 ms	821 ms
`vis_48x32`	1536	517 ms	455 ms	1349 ms
`vis_56x40`	2240	1405 ms	662 ms	2172 ms
`vis_64x36`	2304	1488 ms	681 ms	2278 ms

The runner picks the bucket whose aspect ratio is closest and resizes straight to that grid (no letterbox: padding to a square rewrites line breaks and punctuation). Above ~2.2× linear downscale it falls back to the fp32 ONNX tower, so dense small text is never crushed.

Files

file	what
`vis_*.mlpackage`	vision tower, five fixed grids, fp16 (≈198 MB each)
`prefill_128_2048.mlpackage`	text prefill, 128-token windows, KV 2048, ANE
`decode_1_2048.mlpackage`	decode, one token per call, KV 2048, ANE
`prefill_128_4096.mlpackage`, `decode_1_4096.mlpackage`	same, KV 4096 for long documents
`vision.onnx`, `text_prefill.onnx`	fp32 / int4 CPU fallbacks
`emb_i8.npy`, `emb_i8_scale.npy`, `pos_table.npz`	int8 token embedding, vision pos-embed + rope table
`xocr_ane_full.py`	runner: bucketed ANE vision + ANE prefill + ANE decode
`xocr_ane.py`	decode-on-ANE runner with CPU prefill
`pipeline/`	Core ML graph definitions and build script

Usage

pip install coremltools numpy pillow transformers jinja2 onnxruntime
git clone https://huggingface.co/androidli/Xiaomi-OCR-0-ane
cd Xiaomi-OCR-0-ane
python xocr_ane_full.py --image examples/receipt.jpg --max-new 256

from xocr_ane_full import FullANE
m = FullANE(kv=2048, w=128)
text, tokens, stats = m("page.png", "Extract the text in the image.", 512)
print(stats)   # per-stage timings + which vision bucket was used

Implementation notes (the parts that cost days)

Fixed KV + shift-append. ANE has no dynamic shapes: the 6 full-attention layers keep [1,2,KV,256] and shift one token in per step; absolute positions live in a host-built mask indexed by keys appended so far.
No in-place ops, no aten::Int. Core ML rejects aten::copy_ and every shape-derived int, so the mrope reassembly, the single-token gated-delta update and both block forwards are re-expressed with pure tensor ops and python ints (weights are the original modules).
Core ML has no triangular solve. The chunked gated-delta rule (prefill) needs one; the stable substitute is a blocked inverse (16×16 blocks, T^16 = 0, 3 squarings ≈ 6 matmuls). Full-64 doubling overflows fp32 (L^32 → 2.9e30 → NaN); unrolled row substitution is correct but ~10× larger as a graph.
Converter quirk: converting this graph without extra dummy outputs makes Core ML emit logits that drift by ~40 (CPU and ANE agree with each other, so it is the converter, not the engine). The packages therefore return six unused intermediate tensors — ignore them.

Limits

ANE vision ceiling ≈ 2304 patches; bigger pages are resized into the closest bucket (text verified identical on a 5916-patch receipt) or fall back to ONNX when the downscale is too severe.
M1 Pro and M2 measure the same on the ANE decode (39.7 vs 40.3 ms/step): the graph is operator-dispatch bound, so a newer chip does not help.
The reference runner is Python; for production drive the .mlpackage files from Swift and keep the 48 state tensors in MLState.

Attribution

Base weights: xiaomi-research/Xiaomi-OCR-0, Apache-2.0. Base repo untouched. Nothing fine-tuned.

Made by @greatYiTin.

加载中…