TigreGotico/audiosronnx-mpsenet

audiosronnx-mpsenet

ONNX exports of MP-SENet (Lu et al., MIT) for the mpsenet denoise engine in audiosronnx.

MP-SENet predicts the magnitude and phase spectra in parallel, rather than applying a magnitude mask and reusing the noisy phase. At 2.26 M parameters it is the smallest full spectral model in the library.

amp[1, 201, T], pha[1, 201, T]  ->  amp_g[1, 201, T], pha_g[1, 201, T]

The STFT, the 0.3 power-compression of magnitudes, and the utterance RMS normalisation all run in numpy inside audiosronnx, so inference is onnxruntime-only.

The two checkpoints are not interchangeable. Measured SNR gain on speech corrupted with broadband Gaussian noise, at 19 / 11 / 5 dB input SNR:

dns is the default. vb is the checkpoint the published VoiceBank+DEMAND PESQ figures come from, and is the one to use when reproducing those numbers — but it generalises poorly to noise unlike its training set.

Export note

The model uses stock torch.nn.MultiheadAttention, which specialises the sequence length under both the TorchScript and dynamo exporters; the resulting graph fails outright at any other length. These graphs were exported after swapping in a shape-dynamic attention that reuses the same weights and reproduces the original bit-identically (max abs err 0.0). Verified length-independent from T=50 to T=1500, and end-to-end against the upstream pipeline at correlation 0.99998752.

Usage

from audiosronnx import load_denoise

dn = load_denoise("mpsenet")                # dns checkpoint
dn = load_denoise("mpsenet", model="vb")    # VoiceBank+DEMAND checkpoint
clean, rate = dn.denoise("noisy.wav")       # -> (float32 mono, 16000)

License

MIT, inherited from the upstream MP-SENet release.