androidli/Xiaomi-OCR-0-ane

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:

1. 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.
2. 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.
3. 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.