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 (it reproduces the shipped .tflite byte-for-byte).
• A documented Android (Kotlin) integration recipe below — one CompiledModel.create() call that picks NPU → GPU → CPU by itself.

Not int8. base_model_relation: quantized is HF's closest label; these are fp16 precision + format (ONNX→LiteRT) conversions of the original UVR MDX-Net weights, not integer-quantized models.

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

One call. LiteRT tries NPU, then GPU, then CPU, and keeps whatever works.

Gradle

dependencies {
    implementation("com.google.ai.edge.litert:litert:2.1.6")   // GPU accelerator .so ships inside this AAR
}
android {
    androidResources { noCompress += "tflite" }                // no inflate when copying the asset out
    // ABIs: arm64-v8a (+ x86_64 for the emulator)
}

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

val env = Environment.create(BuiltinNpuAcceleratorProvider(context))
val model = CompiledModel.create(
    tflitePath,   // copy the asset to filesDir once, pass that path — see gotchas
    CompiledModel.Options(Accelerator.NPU, Accelerator.GPU, Accelerator.CPU).apply {
        gpuOptions = CompiledModel.GpuOptions(precision = CompiledModel.GpuOptions.Precision.FP16)
        cpuOptions = CompiledModel.CpuOptions(numThreads = bigCoreCount)   // defaults to 1 — always set it
    },
    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()

Run one warm-up inference on a zeroed input, and retry with Options.CPU if it throws — once the model is compiled LiteRT won't fall back for you.

Gotchas

• numThreads defaults to 1. The CPU rung is XNNPACK, fp16 flag and all, so it's a fine floor — but only if you set the thread count.
• Pass a file path, not context.assets: the assets overload copies the whole model onto the heap (~128 MB peak for Voc FT).
• Nothing tells you which accelerator you actually got. Grep logcat for Replacing N out of M node(s) with delegate.

NPU is work, not a free rung: only SM8550/8650/8750/8850, listed MediaTek SoCs and Tensor G3–G6 qualify, and you ship the Qualcomm/MediaTek libraries yourself (20–110 MB per device, install-time Play Feature Delivery). Untested here; leaving Accelerator.NPU in the set is harmless on everything else.

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. That lone v4 op therefore fails GPU-only compilation outright — and with a CPU rung in the set, silently strands all five decoder upsamplers on the CPU.

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 costs one word in the accelerator set (Accelerator.NPU), but it is hardware-gated (SM8550/8650/8750/8850, listed MediaTek SoCs, Tensor G3–G6) and needs ~20–110 MB of vendor libraries shipped in the APK. Deferred/untested — see above.

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.