litert-community/CMGAN-LiteRT

CMGAN — LiteRT (CompiledModel GPU) speech enhancement / noise suppression

CMGAN (TASLP 2024, MIT, trained on VoiceBank-DEMAND) re-authored to a GPU-native LiteRT .tflite: a 1.83 M-param dual-path conformer (4× time+freq blocks) that denoises 2 s 16 kHz chunks. The STFT (hamming-DFT as one Conv1d) and the mag^0.3 power compression run inside the graph; the host does only reflect-padding, un-compression, inverse STFT and overlap-add. FP16, 4.2 MB.

Noisy input → enhanced output (on-device, +7.2 dB SI-SNR on this clip). Waveforms from the fp16 model.

Verified on a Pixel 8a (Tensor G3): 1 651 / 1 651 nodes LITERT_CL (full residency, 1 partition), ~20 ms per 2 s chunk (RTF ≈ 0.01), SI-SNR +7.2 dB on a 6.6 dB noisy sample (PyTorch reference +9.6 dB), device-vs-PyTorch waveform corr 0.997. fp16 tflite-vs-PyTorch corr 0.999999 (desktop).

I/O

• Input [1, 32400] float32 — a 32 000-sample (2 s @ 16 kHz) mono chunk, RMS-normalized (x·c with c = sqrt(N / Σx²), un-normalize the output by /c) and reflect-padded by 200 samples on both sides (torch.stft(center=True) equivalent).
• Outputs [1, 1, 321, 201] × 2 (real, imag) — the enhanced compressed complex spectrogram (mag^0.3 domain, layout [T, F]). Un-compress by scaling (r, i) by (r²+i²)^(7/6), then iSTFT (n_fft 400 / hop 100 / periodic hamming, trim the 200-sample center pad) and overlap-add chunks.

Minimal usage

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

SR, NFFT, HOP, S = 16000, 400, 100, 32000
wav, _ = sf.read("noisy.wav", dtype="float32")        # mono 16 kHz
x = np.zeros(S, np.float32); n = min(len(wav), S); x[:n] = wav[:n]
c = np.sqrt(S / (x @ x + 1e-12)); x *= c              # RMS normalize
xp = np.concatenate([x[NFFT//2:0:-1], x, x[-2:-NFFT//2-2:-1]])  # reflect pad -> [32400]

it = Interpreter(model_path="cmgan_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], xp[None]); it.invoke()
r, i = (torch.tensor(it.get_tensor(o["index"])) for o in
        sorted(it.get_output_details(), key=lambda o: o["index"]))   # [1,1,321,201] compressed

m2 = (r * r + i * i).clamp_min(1e-12) ** (7.0 / 6.0)                 # mag^0.3 -> mag
spec = torch.complex(r * m2, i * m2)[0, 0].T                         # [201, 321]
den = torch.istft(spec, NFFT, HOP, window=torch.hamming_window(NFFT), length=S)
sf.write("denoised.wav", (den.numpy() / c), SR)

Kotlin (Android, LiteRT CompiledModel GPU)

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

// 2 s chunk, RMS-normalized (x*c), reflect-padded by 200 both sides -> [32400]
inBuf[0].writeFloat(paddedChunk)
model.run(inBuf, outBuf)
val real = outBuf[0].readFloat()   // [321 * 201] compressed real, t-major (t*201 + f)
val imag = outBuf[1].readFloat()
// host: scale (r,i) by (r*r+i*i)^(7/6), then iSTFT (n_fft 400 / hop 100 / periodic hamming,
// trim the 200-sample center pad), overlap-add chunks, divide by c —
// see NoiseSuppressor.kt in the speech_enhancement LiteRT sample.

Conversion (numerically-equivalent re-authoring)

The phase path cancels algebraically (mask·mag·cos(∠x) ≡ mask·x_r) — no atan2/cos/sin in the graph. Shaw relative positional embedding (an Embedding lookup = GATHER) is baked to a constant for the fixed chunk and applied as a 2D FULLY_CONNECTED plus a pad/reshape skew realignment. The conformer's folded batches become batch-1 4D tensors (channel-LayerNorm per position, Linears as 1×1 convs, depthwise Conv1d as (1,k) Conv2d, heads folded into the 3D-BMM batch with 1/√d in Q). mag^0.3 → exp(0.3·ln(·)) (POW is banned); SPConvTranspose2d's 5-D view → an exact 4D reshape chain; InstanceNorm → safe spatial norm; eval-mode BatchNorm → constant scale/shift; all norm eps ≥ 1e-4 (fp16 min-normal on the GPU delegate); no dim-1 broadcast multiplies.

Upstream

• Code + weights: ruizhecao96/CMGAN (MIT)
• Please cite Cao et al., CMGAN: Conformer-Based Metric-GAN for Monaural Speech Enhancement (Interspeech 2022 / TASLP 2024) when you use this model.

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) — cmgan_fp16.tflite	GPU (OpenCL)	—	did not run
TFLite benchmark_model (TfLiteGpuDelegateV2) — cmgan_fp16.tflite	GPU (OpenCL)	352 / 1651	19926.6 ms
TFLite benchmark_model — cmgan_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.

Note that the GPU does not take the whole graph here (352 / 1651 in cmgan_fp16.tflite); the remainder runs on the CPU and the split costs a per-partition round trip.