OmniVoice .tflite Conversion Guide

Prerequisites

x86 Linux machine (litert-torch requires TensorFlow which has no aarch64 wheels for ai-edge-tensorflow)
Python 3.10+
~8GB RAM

Steps

1. Install dependencies

python3 -m venv .venv && source .venv/bin/activate
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip install litert-torch transformers==5.3.0 accelerate soundfile

2. Download from HuggingFace

pip install huggingface_hub
python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('acul3/OmniVoice-LiteRT', local_dir='./omnivoice-litert')
"

3. Convert backbone

import torch
import litert_torch
from export_backbone import OmniVoiceBackbone
from omnivoice import OmniVoice

# Load original model
model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cpu", dtype=torch.float32)

# Extract backbone
backbone = OmniVoiceBackbone(model).eval()

# Sample inputs for tracing
seq_len = 512
input_ids = torch.randint(0, 1024, (1, 8, seq_len), dtype=torch.long)
audio_mask = torch.zeros(1, seq_len, dtype=torch.bool)
audio_mask[:, 200:] = True

# Convert to .tflite
edge_model = litert_torch.convert(backbone, (input_ids, audio_mask))
edge_model.export("omnivoice_backbone.tflite")
print("Backbone exported!")

4. Convert decoder

from export_decoder import OmniVoiceDecoder

decoder = OmniVoiceDecoder(model.audio_tokenizer).eval()
audio_codes = torch.randint(0, 1024, (1, 8, 100), dtype=torch.long)

edge_decoder = litert_torch.convert(decoder, (audio_codes,))
edge_decoder.export("omnivoice_decoder.tflite")
print("Decoder exported!")

5. Verify

# Load and test
import numpy as np
from litert import Interpreter

interp = Interpreter(model_path="omnivoice_backbone.tflite")
interp.allocate_tensors()
print("Backbone loaded OK, input details:", interp.get_input_details())

Critical Notes

Conversion is lossless — identical output to PyTorch
Use FP16 on device — mobile GPUs natively support it
Do NOT use INT8 — quality degrades over 32 diffusion steps
Bidirectional attention is baked into the exported backbone (all-True mask)
Diffusion loop runs in app code (Kotlin), not in the model

Android Integration

// Load models
val backbone = Interpreter(loadModelFile("omnivoice_backbone.tflite"))
val decoder = Interpreter(loadModelFile("omnivoice_decoder.tflite"))

// Enable GPU delegate
val gpuDelegate = GpuDelegate()
val options = Interpreter.Options().addDelegate(gpuDelegate)

// Run diffusion loop
for (step in 0 until numSteps) {
    backbone.run(inputIds, logits)
    // Sample tokens, unmask top-k
}

// Decode to audio
decoder.run(generatedTokens, waveform)