OmniVoice ONNX — bidirectional backbone

An ONNX export of k2-fsa/OmniVoice whose llm_decoder can attend in both directions, which the model requires and the existing export does not do.

Everything except llm_decoder and higgs_decoder is copied unchanged from onnx-community/OmniVoice-Onnx. The Higgs vocoder is that repo's graph converted to fp32, because the fp16 original overflows to NaN on arm64 — see below. The embeddings encoder, heads decoder and tokenizer are untouched.

Why this exists

OmniVoice is a masked diffusion language model. It fills an 8-codebook audio grid by repeatedly committing the slots it is most confident about, so a frame at position 5 must be able to attend to a frame committed at position 50. That is the entire point of ordering commits by confidence.

onnx-community/OmniVoice-Onnx was built with onnxruntime-genai's ModelBuilder, which emits autoregressive decoders: 28 fused com.microsoft.GroupQueryAttention nodes, whose documented contract is causal or local attention only. It takes no arbitrary mask and no attention bias, so the causality cannot be patched out — it lives inside the operator.

The effect is measurable in about ten lines. Run a twelve-token sequence, change only the last token, run it again:

published export:  earlier positions max|diff| = 0.000000e+00   -> causal
this export:       earlier positions max|diff| = 7.312805e+00   -> bidirectional

Bit-identical. Later commitments can never inform earlier ones, the unmasking loop collapses onto a handful of repeated codes, and the output is a buzz rather than speech — 42 distinct codes across 1024 slots, in a faithful reimplementation of the repo's own inference.py.

This is not caught by that repo's eval.py, which validates the embeddings encoder and the heads decoder against PyTorch and describes the LLM as "a black-box genai model". The full chain was never run.

What changed

The Qwen3 backbone was re-exported with torch.onnx.export (opset 20) taking the attention mask as a real 4-D input, which is what transformers passes through verbatim instead of building a causal mask from. Then quantised back to 4 bits.

There are no cache inputs because unmasking runs a full-sequence forward every step and never reuses a cache — the upstream handoff notes say the same.

Files

Pick one backbone from int4/ or fp32/ and take all of components/ with it.

So int4 is about 800 MB all in and fp32 about 2.3 GB. Note that 328 MB of components/ is the three cloning encoders; skip them if you only need text-to-speech.

The layout matters. This was first published flat — every graph at the repo root under a unique name, and the full-precision backbone called llm_decoder_bidir.onnx. Tooling decides "one model in several parts" by finding the same graph basename in more than one directory, so a flat repo reads as eight alternative models instead, and picking any one of them gives you an install that cannot run. Keeping llm_decoder.onnx in both int4/ and fp32/ is what makes the two precisions legible as a choice.

The vocoder had to change too

components/higgs_decoder.onnx as published by onnx-community is float16 from its weights to its output — 136 initialisers, 601 activations, and no Cast node anywhere in the graph. That is fine on x86, where ONNX Runtime's CPU provider has no fp16 kernels and quietly wraps the whole graph in fp32 casts. It is fatal on arm64, where the CPU has real ARMv8.2 half-precision arithmetic and float16 stops at 65504: the decoder overflows, and every sample comes back NaN.

Measured on an 8-core arm64-v8a phone, from a grid that was otherwise healthy — 8 codebooks × 50 frames, 9 to 42 distinct codes per codebook:

vocoder frames=50 n=48000 peak=0 mean=0 nonZero=0 nan=48000 head=NaN, NaN, NaN, NaN

Nothing before it was wrong. Embeddings came back at peak 0.42, backbone hidden states at 179, heads logits at 137, not one non-finite value among them, and the bidirectional probe above passed at 8.21. Only the last graph failed — and it failed silently, because NaN has no peak: a naive silence check reports zero and the result is written out as a valid, empty WAV.

The same file is correct on x86. Run as shipped there it produces clean audio (peak=0.3901 rms=0.0556), and the fp32 conversion in this repo agrees with it to max|diff| = 0.000122, which is fp16 rounding. The graph was never wrong; only its declared precision was, and only on the platform most likely to run it.

The conversion is mechanical — every fp16 initialiser, attribute tensor and value-info upcast to fp32 — and costs 43 MB of download for a vocoder that works everywhere rather than on half of the devices that will try.

It replaces the file in components/ rather than arriving as an fp32/ copy, for the layout reason above: a basename appearing in more than one directory is exactly what tells tooling those directories are alternatives. A second higgs_decoder.onnx under fp32/ would stop fp32/ and int4/ being comparable, and the backbone choice would break. The fp16 original remains in this repo's git history.

Results

Dynamic range across 20 ms frames — speech is intermittent, noise is flat. Judged against the PyTorch reference and confirmed by ear.

Quantisation cost nothing measurable.

Two things the graph alone will not fix

The prompt needs framing. inference.py sends bare tokenizer.encode(text). models/omnivoice.py does not — and without <|text_start|> the model has no marker for where the text begins, so it invents a lead-in and then runs out of grid. This is worth about 33 dB:

<|lang_start|>None<|lang_end|><|instruct_start|>None<|instruct_end|><|text_start|>…<|text_end|>

None is upstream's own literal for "nothing specified". <|denoise|> is prepended only when reference-audio tokens are supplied.

The unmasking loop in inference.py is not the model's. It commits all eight codebooks of a frame at once on a linear schedule. The real one, in omnivoice/models/omnivoice.py, commits individual (codebook, frame) slots chosen by topk over a flattened score grid, on a shifted timestep curve (t_shift=0.1, so the first step commits ~4 slots rather than 32), with a subtractive layer penalty (layer_penalty_factor=5.0) and classifier-free guidance (guidance_scale=2.0) against an unconditional branch that sees the grid with no text.

Guidance is not optional here: with this export, turning it off collapses the output to 0.9 dB. It only makes sense once the unconditional branch can see the whole grid, which is exactly what the causal export prevented.

Usage

import numpy as np, onnxruntime as ort

llm = ort.InferenceSession("int4/llm_decoder.onnx", providers=["CPUExecutionProvider"])

# Full square mask: every position sees every other. This is the input the
# published export does not have.
mask = np.ones((1, 1, S, S), dtype=bool)
hidden = llm.run(["hidden_states"], {
    "inputs_embeds": embeds.astype(np.float32),   # [1, S, 1024]
    "attention_mask": mask,
})[0]

Ask for outputs by name. audio_heads_decoder returns logits shaped [batch, 8, seq, 1025]; the vocoder takes codes shaped [8, 1, frames] and returns waveform_24k at 24 kHz, 40 ms per frame.

Licence and credit

Apache-2.0, inherited from the upstream model. All weights are k2-fsa/OmniVoice (a Qwen3-0.6B finetune); the non-backbone graphs come from onnx-community/OmniVoice-Onnx. This repo contributes a corrected export and the measurements above.

Paper: OmniVoice: Towards Omnilingual Zero-Shot Text-to-Speech with Diffusion Language Models