dellusional/OmniVoice-ONNX-bidirectional

OmniVoice backbone — ONNX export with bidirectional attention

An ONNX export of the k2-fsa/OmniVoice language backbone that preserves the model's bidirectional attention, so it can be used for the masked-diffusion decode OmniVoice actually performs.

Why this export exists

OmniVoice generates audio tokens by iterative unmasking: at every step, masked (codebook, frame) cells must attend to the entire sequence — including frames that come after them. Upstream feeds a fully-True [B, 1, L, L] attention matrix and the model applies no causal masking of its own.

The widely used onnx-community/OmniVoice-Onnx LLM graphs (llm_decoder.onnx and cuda/llm_decoder.onnx) were produced with the autoregressive/KV-cache builder: 28× ORT GroupQueryAttention, and the graph's attention_mask input is consumed only by a ReduceSum that yields seqlens_k — a sequence length, not an attention pattern. Attention is therefore causal, and no caller can override it.

Measured with a single-token perturbation probe (change the last position, then look at how much every earlier hidden state moves):

backbone	max abs delta at earlier positions
upstream PyTorch, 4-D all-True mask	17.24 (bidirectional)
explicit causal mask (control)	0.000000
onnx-community/OmniVoice-Onnx LLM	0.000000 (matches the causal control)
this export	18.10 (bidirectional)

Downstream effect: with the causal graph, iterative unmasking collapses to a handful of repeated codes and the vocoder renders noise. Measured on the same 72-frame utterance, unique codes per codebook were 1–9 with the causal graph versus 50–70 for upstream PyTorch; this export gives 53–71.

What is in the graph

• Plain sdpa forward of the Qwen3 backbone, exported with torch.onnx.export(dynamo=True).
• No KV cache. Every diffusion step recomputes the full sequence anyway, so the cache is dead weight — and it was the KV-cache builder that introduced the causal fusion.
• No embed_tokens. inputs_embeds is always supplied by audio_embeddings_encoder.onnx, so the 151676×1024 embedding is omitted (0.441 B parameters here versus 0.596 B for the full module).

Signature

	name	dtype	shape
input	inputs_embeds	float32	[batch, seq, 1024]
input	attention_mask	bool	[batch, 1, seq, seq] — True = attend
output	hidden_states	float32	[batch, seq, 1024]

attention_mask is the real attention matrix. For ordinary single-sequence generation pass all-True; it is not a padding mask.

Files

file	precision	size
llm_backbone_fp32.onnx + .data	float32	~1.76 GB

An fp16 variant is not published yet. Both routes tried so far fail on this graph: onnxconverter-common leaves mixed-dtype MatMul nodes, and torch.onnx.export(dynamo=True) trips a fake-tensor device mismatch in _scaled_dot_product_efficient_attention when tracing on CUDA and a broadcast error when tracing half weights on CPU. The fp32 graph runs correctly on the CUDA EP (verified below), so this is a size/throughput optimisation, not a correctness gap.

Use with the unmodified audio_embeddings_encoder.onnx, audio_heads_decoder.onnx and higgs_decoder.onnx from onnx-community/OmniVoice-Onnx — those three graphs are correct. (Verified: feeding PyTorch reference codes into the pinned higgs_decoder.onnx reproduces upstream audio at lag-aligned correlation +1.03, offset exactly +100 ms, which is upstream's pad_duration.)

Decoding contract

Reproducing upstream requires more than the graph. In brief:

1. Prompt template — <|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|> followed by <|text_start|>{text}<|text_end|>, each tokenized then repeated across all 8 codebooks, then an all-MASK (id 1024) target region. audio_mask is True only over the target region.
2. instruct is a controlled vocabulary, not free text — e.g. female, young adult, moderate pitch. English items: american/australian/british/canadian/chinese/indian/ japanese/korean/portuguese/russian accent, child, teenager, young adult, middle-aged, elderly, female, male, very low/low/moderate/high/very high pitch, whisper.
3. Classifier-free guidance — a second unconditional pass over the audio region alone, combined as log_softmax(c + 2.0·(c − u)), then the MASK id's logit forced to −inf.
4. Per-cell unmasking — each step reveals the top-k individual (codebook, frame) cells by guided max-log-prob, minus 5.0 · codebook_id, plus Gumbel noise at temperature 5.0, with per-step counts from a t_shift = 0.1 timestep schedule over 32 steps.

Verification

• Torch parity on the exported graph: max abs difference 7.63e-05.
• Bidirectionality probe: see the table above.
• Backbone-swap equivalence: running one deterministic decode twice, changing only the backbone (PyTorch vs this export), produced 575 of 576 identical codes with identical loudness. The single differing cell is a tie-break at fp32 rounding scale.
• End-to-end synthesis through the three pinned graphs produces finite audio with speech-like dynamics (peak ≤ 1.05, rms ≈ 0.12) and reference-band code diversity.
• CUDA EP (onnxruntime-gpu 1.23.2, CUDA 12): loads and runs, 34.1 ms per forward at sequence length 96, zero non-finite outputs, bidirectionality preserved.

Export script

llm = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cpu", dtype=torch.float32).llm.eval()
llm.set_input_embeddings(torch.nn.Embedding(1, llm.config.hidden_size))  # inputs_embeds is always supplied

class Backbone(torch.nn.Module):
    def forward(self, inputs_embeds, attention_mask):
        return self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, return_dict=True)[0]

torch.onnx.export(
    Backbone(llm).eval(), (embeds, mask), "llm_backbone_fp32.onnx",
    input_names=["inputs_embeds", "attention_mask"], output_names=["hidden_states"],
    dynamic_shapes={"inputs_embeds": {0: batch, 1: seq},
                    "attention_mask": {0: batch, 2: seq, 3: seq}},
    dynamo=True, external_data=True, opset_version=18,
)

License and attribution

CC-BY-NC 4.0 — non-commercial use only.

Upstream states it directly on the model card: "Our code is released under the Apache 2.0 License. The pre-trained model is licensed under the CC-BY-NC due to constraints from its training data (e.g., Emilia)." The Apache-2.0 grant covers the k2-fsa code; it does not cover these weights. This repository contains weights, so CC-BY-NC governs, and it propagates to anything derived from them.

This is a derived work: a re-export of the upstream weights into ONNX with no fine-tuning and no modification of any parameter value — only the graph's attention masking and I/O contract differ. All model credit belongs to the k2-fsa authors.

Upstream's disclaimer applies here in full: users are prohibited from using this model for unauthorized voice cloning, voice impersonation, fraud, scams, or any other illegal or unethical activity.

Related components, under different terms

The companion graphs named above come from onnx-community/OmniVoice-Onnx and are not redistributed here. Note that higgs_decoder.onnx derives from upstream's audio_tokenizer/, which carries the Boson Higgs Audio 2 Community License (itself built on the Meta Llama 3 Community License) — a separate agreement from CC-BY-NC. Check both before using the full pipeline.