colSmol-256M-dynamic-onnx

An ONNX export of vidore/colSmol-256M that accepts a dynamic tile count.

Why this exists

The existing community export, onnx-community/colSmol-256M-ONNX, is traced at exactly 13 tiles and rejects any other count. Idefics3 tiles an image according to its aspect ratio, so the tile count is a property of the input, not a constant:

aspect ratio	tiles
16:10 (typical desktop)	7
landscape	13
5:4 / square	17

A 13-tile-only export therefore cannot process a 16:10 screenshot at all. This export makes the tile axis dynamic, so one model handles every aspect ratio.

What changed

Two places in Idefics3 use data-dependent shapes, which is what bakes the tile count into a trace. Both are replaced with shape-static equivalents:

1. Idefics3Model.get_image_features filters "padding" tiles with boolean-mask indexing (pixel_values[real_images_inds]). The result's size depends on the data, so tracing freezes it. The filter is dropped — callers must pass only real tiles, which is this export's documented precondition. The per-patch attention mask is also rebuilt with a reshape-and-sum instead of two unfold calls, which ONNX cannot export once the leading dim is dynamic.

2. Idefics3VisionEmbeddings.forward does a masked assignment (position_ids[mask] = pos_ids[mask]). Replaced with torch.where, which is numerically identical — masked-out positions keep the zero fill — but static.

Precondition: pass only real tiles. This export has no padding-tile filter, so padding tiles would be embedded as if they were content.

Inputs and outputs

inputs : input_ids, attention_mask, pixel_values, pixel_attention_mask
outputs: embeddings

Dynamic axes: batch and sequence on the token inputs, batch and tiles on the pixel inputs. Output is (batch, seq, 128) where seq = tiles * 64 + 2.

Validation

Verified against eager PyTorch with the same patches applied:

tiles	output dims	expected (n*64+2)
7	(1, 450, 128)	450
13	(1, 834, 128)	834
5	(1, 322, 128)	322
1	(1, 66, 128)	66

Numerical parity at 7 tiles: max abs diff 5.119e-06, min cosine 1.000000.

Steady-state CPU inference is roughly 5.0s per frame at 7 tiles (CPUExecutionProvider, fp32).

Files

file	source
model.onnx	exported here from vidore/colSmol-256M
tokenizer.json, tokenizer_config.json, preprocessor_config.json	onnx-community/colSmol-256M-ONNX @ 12022f06
config.json	vidore/colSmol-256M @ a59110fd

The companion configs are bundled so they cannot drift out of sync with the weights they were traced against.

Reproducing

Built with scripts/export-colsmol.py from DeskRAG, CPU/fp32, opset 17, traced at 7 tiles:

package	version
torch	2.11.0
transformers	5.14.1
onnx	1.22.0
colpali-engine	0.3.17

The monkeypatches target version-sensitive transformers internals, so those pins matter when regenerating.

License

MIT, following vidore/colSmol-256M, whose adapters are MIT over an Apache-2.0 SmolVLM backbone. Both permit redistribution of derivative works.