unirig-skeleton-onnx-webgpu

UniRig skeleton stage, quantized ONNX for the browser

UniRig ("One Model to Rig Them All", SIGGRAPH 2025, VAST-AI-Research) predicts a skeleton from raw mesh geometry, no template and no markers. This is the skeleton stage quantized to int8 so it runs entirely inside a web browser on WebGPU.

1.44 GB down to 455 MB, with the same skeleton coming out.

Try it without downloading anything

Make a 3D model at freegen.ai

Files

Inputs and outputs

• encoder: vertices [1,N,3], normals [1,N,3] float32, N up to 65536, surface-sampled, centred, and scaled so the longest side spans -1 to 1, with +Y up, to latents [1,1024,1024].
• embed: input_ids [1,S] int64 to token_embeds [1,S,1024].
• decoder: inputs_embeds [1,S,1024] and past.{0..23}.{key,value} [1,16,P,64] to logits [1,S,267] and present.{0..23}.{key,value}.

How it is driven

latents = encoder(points, normals)
inputs_embeds = concat([latents, embed([bos=257, cls=266])])
loop:
    logits, present = decoder(inputs_embeds, past)
    token = argmax over ONLY the tokens the grammar allows
    inputs_embeds = embed([token]); past = present

The class token is easy to miss and is not optional.

The box is -1 to 1, and this matters more than it sounds

"A centred unit box" is ambiguous and the wrong reading does not give a scaled rig, it gives no rig. At -0.5 to 0.5 the decode never emits its end token and runs to the safety cap, producing a hundred bones in a line. The same cloud scaled to -1 to 1 terminates normally at twenty two. Measured both ways on the same points.

The grammar mask is load-bearing

256 coordinate bins over [-1,1], undiscretize(t) = (t+0.5)/256*2-1; vocab 267 with branch=256, bos=257, eos=258, cls=266. Bones are [branch?, parent xyz, joint xyz], so six coordinates each; eos and branch are legal only on a six-coordinate boundary.

This is not defensive. On the first step of a real decode the unmasked argmax is token 260, an unused slot in the vocabulary. An unconstrained greedy decode produces garbage immediately.

Match a parent to the nearest earlier joint rather than to an exact bin; real decodes routinely miss the bin the parent was written to.

What was changed, and what was measured

• Gemm rewritten to MatMul + Add before quantizing. Two thirds of the decoder's weight, 806 MB of 1223 MB, is in Gemm nodes, which MatMulNBitsQuantizer does not touch. Without this the decoder only gets to 933 MB. Only the exact form these graphs use is rewritten (alpha=beta=1, no transA, transB set, initializer weight).
• The output projection stays in float32. Every other weight feeds a residual stream that absorbs a small error. This one produces the logits an argmax is taken over, so an error becomes a different token, and the next step is conditioned on it. Quantizing it too put the first divergence at token 81 of 167 and turned a 27 bone rig into a 29 bone one. Excluding it costs about a megabyte.
• int8, not int4. The decode is autoregressive; a token chosen wrongly at step nine is not a slightly worse rig, it is a different rig from step nine on.
• Block 16 for the decoder, 32 for the encoder. The two graphs are not equally forgiving. The encoder at block 32 is exact: behind a float32 decoder it reproduces all 112 tokens of the reference decode. The decoder at block 32 diverges on the first token and turns an eighteen bone rig into a twenty two bone one; at block 16 it agrees on 111 of 112. That is 38 MB for a correct rig.

Verification

scripts/unirig-verify.py runs the whole chain on a real reconstruction rather than on noise, and compares the quantized pack against the original on the only thing that matters: 111 of 112 tokens identical, diverging only on the last, the same 18 bones, the hierarchy identical, and joints within 1.00 quantisation bins (one bin is the finest distance the model can express). A different hierarchy, or joints moving more than a bin and a half, fails the build.

Note on the upstream ONNX export

The excellent export this is derived from is by fernandotonon for QtMeshEditor. If you go looking for it: the graphs that match the documented contract are in fernandotonon/QtMeshEditor-models under unirig/. The standalone QtMeshEditor-unirig-onnx repository currently mirrors a generic input_ids export of the decoder, which cannot be conditioned on a mesh at all.

Licence

MIT, from the upstream code and weights. Trained on Articulation-XL2.0 (CC-BY-4.0) — credit VAST-AI-Research. Permits commercial use.