arlenecc/OvisOCR3-ONNX-int4

OvisOCR3 — ONNX INT4（单文件版）

整个模型打包成 一个 ovisocr3_int4_single.onnx（1.16 GiB）， 视觉编码 / 批量 prefill / 单步解码 / 嵌入查表全部在这一个文件里。

为什么是多路复用的一张图

模型里有动态形状的部分（图像 patch 数、past 长度），但这些只是张量维度， 真正的问题是 4 段计算各自的输入集合不同，而一个 ONNX 模型只有一个 graph。 做法是把 4 张子图用 If(mode) 包起来放进同一张图的分支里：

input mode ─┬─ 0 → 视觉编码 branch   (vision)
            ├─ 1 → 批量 prefill      (32 token/次)
            ├─ 2 → 单步解码          (内含 INT4 嵌入查表)
            └─ 3 → 嵌入查表          (GatherBlockQuantized)

• 每个分支保留自己的 initializer（MatMulNBits / GatherBlockQuantized 权重）
• 分支内部张量全部加分支前缀，避免与外层输入名冲突
• 嵌入权重提升到主图共享，避免 step 分支内联时重复 155 MB
• 分支图为 0 输入，通过隐式捕获引用外层（主图）的输入张量 —— opset 21 的 If 只允许 condition 一个显式输入，这是唯一可行的写法

mode=2（解码）分支把嵌入 Gather 内联了，所以解码每步不再单独跑一次查表图， 只传 input_ids 即可。

文件

文件	大小	说明
onnx/ovisocr3_int4_single.onnx	1.16 GiB	全部内容，单文件无外挂数据
run_ocr.py		推理 CLI
single_model.py		单文件版推理封装（OvisSingle）
ort_runner.py / preprocess.py		基础推理封装与预处理

对比其余两个多文件版本：

版本	ONNX 文件数	总体积
dist（嵌入 bf16 npy）	3 onnx + 1 npy	1.49 GB
dist-int4-full（嵌入也 INT4）	4 onnx	1.16 GB
dist-single（本目录）	1 onnx	1.16 GB

运行

python run_ocr.py --image /path/to/page.png --prompt plain_ocr --max-new 256 --onnx-dir ./onnx

默认自动探测 onnx/ovisocr3_int4_single.onnx；--multi 可强制走多文件版。 需要 onnxruntime ≥ 1.20（用到 MatMulNBits、GatherBlockQuantized）。

图的输入输出

输入（每次调用都要传，未使用的给空数组即可）：

名字	形状	dtype	用途
mode	标量	int64	0/1/2/3
pixel_values	[N, 1536]	float	展开后的 patch 像素
pos_embeds	[N, 768]	float	插值好的 2D 位置嵌入
position_ids	[N, 2]	int64	视觉 (H, W) 位置
input_ids	[C]	int64	查表用的 token id
inputs_embeds	[1, C, 1024]	float	prefill 输入（已注入视觉 token）
cos / sin	[1, C, 64]	float	mrope 预计算
past_key / past_value	[6, 1, 2, P, 256]	float	full attention KV cache
conv_state	[18, 1, 6144, 3]	float	GDN 卷积状态
rec_state	[18, 1, 16, 128, 128]	float	GDN 递归状态

输出：vision_tokens / emb_out / logits / present_key / present_value / conv_state_out / rec_state_out（未走到的分支返回零尺寸张量或状态直传）。

正确性

22_check_single.py 逐 mode 与原来的独立模型对比， 4 个分支的输出全部 max_abs = 0.000e+00（逐位相同）； 端到端 OCR 文本与多文件版完全一致。

性能（i5-8259U 4 线程，638 token 提示 + 64 生成）

	多文件版	单文件版
ONNX 文件数	4	1
加载耗时	13.0 s	29 s
prefill	8.5–11.9 s	9.4–9.9 s
单步解码	69–73 ms	73–76 ms
解码速度	13.7–14.4 tok/s	13.2–13.7 tok/s（**约 -5%**）

代价主要是：If 分支内的子图拿不到和顶层图同等程度的优化， 以及所有分支权重常驻同一个 session。换来的是只有 1 个文件。

已知限制：

• 预处理（tokenizer / 图像处理器 / mrope 位置 id）仍依赖 transformers
• 仅支持 batch=1
• mode=2 的 step 分支内联了嵌入查表，因此不能直接喂图像 token 的嵌入 （驱动里已加保护，遇到会直接报错）。实际使用时图像 token 只出现在 prefill 段