TIGER-DnR β€” LiteRT (CompiledModel GPU) cinematic sound separation

on-device result

Mixture input β†’ separated Dialogue / Sound effects / Music stems, from the on-device model (waveforms).

TIGER (ICASSP 2025) DnR cinematic sound separation re-authored to GPU-native LiteRT .tflite: split a clip into Dialogue / Sound effects / Music stems fully on-device. Three sibling band-split TIGER graphs (~1.41 M params each, trained on the openly-built DnR dataset by the upstream authors) each process a 12.06 s, 44.1 kHz mono chunk; per DnR convention each graph contributes one stem (dialog = source 2 of tiger_dialog, effect = source 1 of tiger_effect, music = source 0 of tiger_music).

Verified on a Pixel 8a (Tensor G3): full LITERT_CL residency β€” 23 974 / 23 974 nodes, 1 partition per graph, GPU output vs PyTorch waveform corr 0.99987, ~4.5 s per chunk per graph. FP16, 16.1 MB per graph.

Files

file stem it contributes
tiger_dialog_fp16.tflite Dialogue (output source index 2)
tiger_effect_fp16.tflite Sound effects (output source index 1)
tiger_music_fp16.tflite Music (output source index 0)

I/O

  • Input [1, 534016] float32 β€” a 531 968-sample (12.06 s @ 44.1 kHz) mono chunk, reflect-padded by 1024 samples on both sides by the caller (torch.stft(center=True) equivalent). The STFT (win 2048 / hop 512, periodic Hann) runs inside the graph as a windowed-DFT Conv1d.
  • Outputs [1, 3, 1025, 1040] Γ— 2 (real, imag) β€” separated one-sided complex spectrograms for the 3 sources. Host runs iSTFT (win 2048 / hop 512, Hann, trim 1024 center pad) and overlap-adds chunks (reference inference uses a 12 s window with 4–10 s hop, averaging overlaps).

Minimal usage (Python, desktop parity check)

import numpy as np, soundfile as sf, torch
from ai_edge_litert.interpreter import Interpreter

SR, WIN, HOP, T = 44100, 2048, 512, 1040
S = (T - 1) * HOP                                    # 531968 samples = 12.06 s chunk

wav, _ = sf.read("mix.wav", dtype="float32")         # mono 44.1 kHz
x = np.zeros(S, np.float32); n = min(len(wav), S); x[:n] = wav[:n]
x = np.concatenate([x[WIN//2:0:-1], x, x[-2:-WIN//2-2:-1]])   # reflect pad (host side)

it = Interpreter(model_path="tiger_dialog_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x[None]); it.invoke()
real, imag = (it.get_tensor(o["index"]) for o in
              sorted(it.get_output_details(), key=lambda o: o["index"]))  # [1,3,1025,1040] x2

spec = torch.complex(torch.tensor(real), torch.tensor(imag))[0, 2]       # source 2 = dialogue
dialog = torch.istft(spec, n_fft=WIN, hop_length=HOP,
                     window=torch.hann_window(WIN), length=S)
sf.write("dialog.wav", dialog.numpy(), SR)
# effects: tiger_effect_fp16.tflite source 1 Β· music: tiger_music_fp16.tflite source 0

Kotlin (Android, LiteRT CompiledModel GPU)

// implementation("com.google.ai.edge.litert:litert:2.1.5")
val model = CompiledModel.create(File(ctx.filesDir, "tiger_dialog_fp16.tflite").absolutePath,
    CompiledModel.Options(Accelerator.GPU), null)
val inBuf = model.createInputBuffers()
val outBuf = model.createOutputBuffers()

// 12.06 s mono 44.1 kHz chunk (531968 samples), reflect-padded by 1024 on both sides -> 534016
inBuf[0].writeFloat(paddedChunk)
model.run(inBuf, outBuf)
val real = outBuf[0].readFloat()   // [3 * 1025 * 1040], source-major: src*1025*1040 + f*1040 + t
val imag = outBuf[1].readFloat()
// dialogue = source index 2; host iSTFT (win 2048, hop 512, periodic Hann, trim 1024 center pad)
// + overlap-add across chunks β€” see Istft.kt in the audio_source_separation LiteRT sample.

Conversion (numerically-equivalent re-authoring)

torch.stft β†’ in-graph DFT-as-Conv1d (host reflect-pad); torch.istft β†’ host; folded-batch Conv1d ((BΒ·T, N, band)) β†’ 4D (1,k)-Conv2d; per-sample GlobLN β†’ per-position chained-single-axis-mean SafeNorm; chunk length chosen so T=1040 is divisible by 16 β†’ adaptive pools become uniform AVERAGE_POOL_2D and nearest resizes exact integer repeats; the non-uniform 57-band axis uses constant averaging/one-hot FULLY_CONNECTED matrices; MHSA β†’ per-head batch-1 3D matmuls with 1/√d folded into Q; PReLU β†’ relu(x) βˆ’ wΒ·relu(βˆ’x); the 6-D mask-head view β†’ static channel slices. Two fp16-on-GPU fixes, both exact-equivalent: norm eps raised to 1e-4 (1e-8/1e-5 underflow to 0 in fp16 β†’ 0/0 NaN on silent bands) and the mask head rewritten without dim-1 broadcast MULs.

Desktop parity: fp16 tflite vs PyTorch waveform corr 0.99991 (per-graph, real 12 s mixture).

Upstream

Please cite the TIGER paper (Xu et al., ICASSP 2025, arXiv:2410.01469) when you use these models.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool β€” 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime Backend Graph on GPU Latency
TFLite benchmark_model (TfLiteGpuDelegateV2) β€” tiger_effect_fp16.tflite GPU (OpenCL) 23974 / 23974 4974.9 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) β€” tiger_music_fp16.tflite GPU (OpenCL) 23974 / 23974 5074.1 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) β€” tiger_dialog_fp16.tflite GPU (OpenCL) 23974 / 23974 5140.1 ms
TFLite benchmark_model β€” tiger_effect_fp16.tflite CPU (XNNPACK, 4 threads) β€” XNNPACK declined the graph
TFLite benchmark_model β€” tiger_music_fp16.tflite CPU (XNNPACK, 4 threads) β€” XNNPACK declined the graph
TFLite benchmark_model β€” tiger_dialog_fp16.tflite CPU (XNNPACK, 4 threads) β€” XNNPACK declined the graph

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

XNNPACK declines these fp16 graphs β€” it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors β€” so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20Γ— slower than the GPU on models of this size and would not represent CPU inference anyone would ship.

Downloads last month
311
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Model tree for litert-community/TIGER-DnR-LiteRT

Finetuned
(1)
this model

Collection including litert-community/TIGER-DnR-LiteRT

Paper for litert-community/TIGER-DnR-LiteRT