gyoom-sa/UVR-MDX-LiteRT

UVR MDX-Net → LiteRT (on-device vocal separation)

Two UVR MDX-Net vocal separators, converted to LiteRT (.tflite) as fp16 and made to delegate whole to a mobile GPU — the models that power on-device vocal/instrumental separation in the MusicStemSeparation Android app. This repo ships the models plus everything needed to run and reproduce them:

• separate.py — a self-contained, runnable Python reference (numpy STFT + LiteRT). WAV in → vocals + instrumental WAV out.
• export/export_mdx_litert.py — the single, deterministic ONNX→LiteRT-fp16 export script
• A documented Android (Kotlin) integration recipe below — the real CPU→GPU accelerator ladder.

Models

Model	Params	dim_f	n_fft	hop	Size	Notes
UVR_MDXNET_9482.fp16acc.tflite	7.4 M	2048	4096	1024	~28 MB	general, fast (default)
UVR-MDX-NET-Voc_FT.fp16acc.tflite	16.7 M	3072	6144	1024	~64 MB	vocal fine-tune, higher quality

Both run a native 256-frame segment (dim_t = 256, ≈ 5.9 s @ 44.1 kHz). The model predicts the vocal spectrogram; the instrumental is the free residual mix − vocals.

Quickstart

pip install -r requirements.txt          # numpy + ai-edge-litert (WAV IO is stdlib)

# input must be WAV — convert first if needed:
ffmpeg -i example.mp3 example.wav

python separate.py example.wav                    # 9482 (default)  → example_vocals.wav + example_instrumental.wav
python separate.py example.wav --model voc_ft     # higher-quality model
python separate.py example.wav --model voc_ft --denoise --out-dir out/

separate.py runs the model on the CPU (via ai-edge-litert/XNNPACK) — a correct, portable reference. The fp16 speedup and the GPU/NPU path are Android-only (see below).

Model I/O contract (the one thing to get right)

input = OUTPUT = float32 tensor  [1, 4, dim_f, 256]   (NCHW)

The 4 channels are complex-as-channels, plane order [L_re, L_im, R_re, R_im] (real plane = 2·c, imag plane = 2·c + 1), flat row-major index ((plane·dim_f) + bin)·256 + frame.

• The STFT keeps bins [0, dim_f) of the n_fft/2 + 1 one-sided bins — i.e. the Nyquist bin is dropped (dim_f = n_fft/2); on the inverse it is zero-padded back.
• Voc FT carries no embedded STFT metadata — derive n_fft = dim_f · 2.
• A wrong packing (swapped re/im planes, a kept Nyquist bin, a stray 1/√n) produces plausible-but-wrong audio, not an error. If a stem sounds off, look here first.

How it works — STFT outside the graph

The .tflite is the learned core only — a convolutional U-Net that maps a spectrogram to the vocal spectrogram. The STFT, chunking, overlap-add and inverse STFT are ordinary CPU code (numpy in separate.py, Kotlin in the app). Per chunk:

STFT (host CPU) → inference (LiteRT: CPU / GPU / NPU) → iSTFT (host CPU) → weighted overlap-add

This is a faithful reimplementation of UVR's mdx.py demix: periodic Hann window, center=True reflect padding, unnormalized (torch.stft(normalized=False)), a fixed 10% crossfade between chunks (set to 0 for the classic hard-tiled concat), and the free instrumental residual. Because the accelerator only ever sees the learned core, the model delegates cleanly and every stage is independently timed.

Android on-device inference (XNNPACK ↔ GPU)

The app picks the accelerator by availability, not by assumption. LiteRT exposes two APIs, and each rung uses the right one:

• XNNPACK (fp16) / built-in CPU → the classic org.tensorflow.lite.Interpreter.
• GPU (and NPU) → the newer com.google.ai.edge.litert.CompiledModel + Accelerator.GPU.

CompiledModel is all-or-nothing — one unsupported op fails the whole compile (that is exactly why the export un-fuses TRANSPOSE_CONV, see below). And whether the GPU delegate accepts a model can't be known without trying (coverage-gated, not accelerator-exists). So: build each rung, run a warm-up inference, and keep the first that survives, falling back GPU → XNNPACK → CPU.

Gradle

dependencies {
    implementation("com.google.ai.edge.litert:litert:2.1.6")   // GPU accelerator .so ships inside this AAR
}
android {
    androidResources { noCompress += "tflite" }                // let LiteRT memory-map the model
    // ABIs: arm64-v8a (+ x86_64 for the emulator)
}

Snippet A — XNNPACK / CPU via Interpreter (fp16 engages automatically on ARMv8.2-FP16 cores)

import org.tensorflow.lite.Interpreter
import java.io.RandomAccessFile
import java.nio.channels.FileChannel

val mapped = RandomAccessFile(tflitePath, "r").use { raf ->
    raf.channel.map(FileChannel.MapMode.READ_ONLY, 0, raf.length())   // mmap the .tflite
}
val interpreter = Interpreter(mapped, Interpreter.Options().apply {
    setNumThreads(bigCoreCount)   // e.g. the device's performance-core count
    setUseXNNPACK(true)           // fp16 kicks in via the model's reduced_precision_support flag
})
// per chunk: input/output are direct float ByteBuffers of 1*4*dim_f*256 floats
interpreter.run(inputBuffer, outputBuffer)

Snippet B — GPU via CompiledModel (fp16 is the GPU's native, ~2× precision; MDX is fp16-safe)

import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel
import com.google.ai.edge.litert.Environment

val env = Environment.create()
val model = CompiledModel.create(
    tflitePath,
    CompiledModel.Options(Accelerator.GPU).apply {
        gpuOptions = CompiledModel.GpuOptions(precision = CompiledModel.GpuOptions.Precision.FP16)
    },
    env,
)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
// per chunk:
inputs[0].writeFloat(chunkFloats)   // 1*4*dim_f*256, [L_re, L_im, R_re, R_im]
model.run(inputs, outputs)
val vocals = outputs[0].readFloat()
// NPU: swap Accelerator.NPU + QualcommOptions(HtpPerformanceMode.SUSTAINED_HIGH_PERFORMANCE);
// needs a vendor runtime (hardware-gated — see "NPU" below).

Snippet C — the fallback ladder (build + warm-up-run each rung; keep the first that survives)

fun createEngine(tflitePath: String): Engine {
    for (backend in listOf(Backend.GPU, Backend.XNNPACK, Backend.CPU)) {
        runCatching {
            val engine = build(backend, tflitePath)   // Snippet A or B
            engine.warmUp()                            // one inference on a zeroed input
            engine
        }.onSuccess { return it }
         .onFailure { Log.w(TAG, "backend $backend failed; falling back", it) }
    }
    error("no backend available")   // the built-in CPU floor is always last and always works
}

The full production version (fallback-reason reporting, the NPU rung, warm-up gating, the MdxEngine seam) lives in the app's tools/mdxlitert/pipeline/MdxLiteRtEngine.kt.

Why LiteRT — the acceleration story

MDX-Net delegates whole. It is a clean single-domain conv U-Net — no attention, no LSTM, no whole-tensor norm — which is exactly the op set XNNPACK, the GPU (ML Drift / OpenCL) and (on capable SoCs) the NPU map end-to-end. With the STFT pulled out of the graph, the accelerator sees only convs, transposed-convs and the TDF fully-connected blocks, so the whole graph delegates in one partition.

fp16 is the ~2× lever and it's safe here. fp16 (not int8) is the quality-safe speedup, but only for a model with no whole-tensor reduction to overflow the fp16 ceiling (65504). MDX has only per-channel BatchNorm, which folds into the conv — nothing overflows:

Model	peak activation	fp16 ceiling	headroom
9482	559	65504	117×
Voc FT	1384	65504	47×

fp16 is requested via the .tflite metadata flag reduced_precision_support = "fp16accfp16" (the fast accumulate value — measured quality-safe for MDX), which lets XNNPACK run the delegated ops in fp16 on ARMv8.2-FP16 cores.

The GPU fix — un-fuse TRANSPOSE_CONV's ReLU

The converter fuses each decoder ConvTranspose + BN + ReLU into one TRANSPOSE_CONV op with a fused ReLU — and a fused activation is precisely what tags the op v4. LiteRT 2.1.6's OpenCL delegate caps TRANSPOSE_CONV at v3, whose semantics have no fused activation. Because CompiledModel is all-or-nothing, that lone v4 op fails the whole GPU compile.

The export un-fuses it: each TRANSPOSE_CONV(+fused ReLU) becomes a legitimately-v3 bias-only TRANSPOSE_CONV followed by a standalone RELU op (which the GPU supports). The whole graph then delegates in one partition and computes the right function. (Simply force-lowering the v4 tag to v3 instead — the earlier attempt — made the GPU accept the graph but silently drop the ReLU on all five decoder upsamplers, producing an audible broadband buzz.) CPU numerics are unchanged.

Reproducing the export

The two .tflite were produced by export/export_mdx_litert.py, and the export is deterministic — re-running it from the source ONNX yields the shipped files byte-for-byte (verified via SHA-256).

# heavy converter env (Python 3.11):
pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cpu
pip install -r export/requirements.txt

python export/export_mdx_litert.py UVR_MDXNET_9482.onnx     UVR_MDXNET_9482.fp16acc.tflite
python export/export_mdx_litert.py UVR-MDX-NET-Voc_FT.onnx  UVR-MDX-NET-Voc_FT.fp16acc.tflite

Six gated steps: onnx2torch → SNR-gate (>100 dB vs ONNX Runtime) → litert_torch export (NCHW, static dim_t=256) → +fp16 metadata → un-fuse TRANSPOSE_CONV ReLU → verify (CPU SNR, output shape, no op above the GPU v3 cap). Source ONNX: the UVR MDX-Net releases from TRvlvr/model_repo, also on the Hub at Politrees/UVR_resources.

onnx2tf does not work here — its per-op NCHW→NHWC guessing breaks the TDF MatMuls. The PyTorch → StableHLO → TFLite route (litert_torch) does no layout guessing and keeps NCHW I/O, so the app's complex-as-channels packing feeds the model directly.

Fidelity & performance

• Fidelity: the fp16 .tflite is gated at > 100 dB CPU SNR vs the ONNX reference; measured ≈ 109 dB (9482) and ≈ 112 dB (Voc FT) after the un-fuse step. Validated by ear on-device.
• Performance (one measured device — Xiaomi Redmi Note 10 Pro, Snapdragon 732G / Adreno 618, Android 13): XNNPACK-fp16 runs the graph at RTF ≈ 1.0 (roughly real-time) as the universal CPU floor; the GPU path (whole-graph LITERT_CL partition) runs clean and faster than XNNPACK after the un-fuse fix. Numbers vary by SoC; treat these as one data point, not a benchmark.
• NPU: the rung is wired (CompiledModel + Accelerator.NPU, QNN / NeuroPilot) but is hardware-gated — LiteRT admits only SM8550+ Snapdragons and listed MediaTek Dimensity SoCs — and is currently deferred/untested.

Attribution & license

Released under the MIT License, matching the upstream UVR MDX-Net models.

• UVR MDX-Net and the 9482 model — the Ultimate Vocal Remover project by Anjok07 and contributors.
• UVR-MDX-NET Voc FT — vocal fine-tune by Kimberley Jensen.
• Source ONNX via TRvlvr/model_repo (Hub mirror: Politrees/UVR_resources).

This repo's contribution: the ONNX → LiteRT fp16 conversion, the whole-graph GPU fix (un-fusing TRANSPOSE_CONV's ReLU), the runnable separate.py reference, and the Android integration recipe. The weights are the original UVR models, format-converted.