colSmol-256M-dynamic-onnx

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:

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:

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

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:

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.