litert-community/TIGER-DnR-LiteRT

TIGER-DnR — LiteRT (CompiledModel GPU) cinematic sound separation

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

• Code: JusperLee/TIGER (MIT)
• Weights: JusperLee/TIGER-DnR (Apache-2.0)

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