Resemble Enhance (ONNX)

ONNX export of Resemble Enhance (Resemble AI, MIT), the speech restoration model. The networks come from the official enhancer_stage2 checkpoint (ResembleAI/resemble-enhance); the pipeline (STFT, mel features, denoiser post-processing, flow-matching solver, crossfade) is the same as the earlier LiteRT version, implemented in inference.py on ONNX Runtime.

Five networks run in ONNX Runtime, and the signal processing runs in numpy:

Network Input β†’ output Role
denoiser-net features (1, 3, 841, 105) β†’ (1, 3, 841, 105) predicts the denoiser mask and phase residual
encoder mel (1, 128, 105) β†’ latent (1, 64, 105) encodes the mel spectrogram
flow-step state, condition, time β†’ velocity (1, 64, 105) one evaluation of the flow-matching velocity field
decoder latent (1, 64, 105) β†’ conditioning (1, 160, 105) decodes the flow output
vocoder conditioning, noise (1, 128, 115) β†’ waveform (1, 44100) UnivNet vocoder, one 1-second window

Files

Variant Files Size What is quantized
fp32 *-float32.onnx ~1.4 GB nothing: the reference export
fp16 *-fp16.onnx ~0.7 GB all networks in float16 (inputs and outputs stay float32)
int8 denoiser-net, encoder, flow-step, decoder as *-int8-dynamic.onnx; vocoder-fp16.onnx ~0.5 GB INT8 dynamic weights for the four small networks; the vocoder stays fp16 (INT8 vocoder weights degraded quality)

mel-filters.f32 is the mel filter bank (1025 Γ— 128, float32 little-endian), from the original mel front end.

Accuracy

Each export matches PyTorch. With random inputs, the fp32 ONNX networks differ from the PyTorch modules by at most 5.6e-4 (the denoiser, whose output reaches 17) and by about 1e-6 to 5e-6 for the others.

The full pipeline was compared on one 11-second speech clip (JFK, 44.1 kHz), using the same fixed seeds for every variant:

Variant compared with fp32 Waveform SNR Log-spectral distance
fp16 27.3 dB 0.061 (β‰ˆ 0.6 dB)
int8 (this repository) 4.5 dB 0.26 (β‰ˆ 2.6 dB)

The log-spectral distance is the mean over frames of the RMS difference of log power spectra (log10 units; multiply by 10 for dB). Waveform SNR is strict for a generative vocoder: a small change in the conditioning can shift phase and lower SNR even when the sound stays similar. Listen to the outputs before relying on int8. fp16 is the variant we recommend when size allows.

The per-network differences are larger on real speech than on random inputs. For example, the INT8 denoiser differs from fp32 by about 50% (relative to its largest value) on the clip, and the differences grow through the 16 flow steps into the vocoder.

Usage

pip install onnxruntime numpy scipy soundfile huggingface_hub
python inference.py input.wav output.wav --dtype fp16

The output is 44.1 kHz mono 16-bit WAV. Processing runs in 2-second windows with 50% overlap and a Hann crossfade. Each window's noise comes from a fixed seed, so the same input gives the same output.

How it was made

  • Source: the enhancer_stage2 run from ResembleAI/resemble-enhance (MIT).
  • conversion/export_onnx.py exports each network with torch.onnx.export (opset 17; the vocoder uses the dynamo exporter at opset 18). Before export it applies the same rewrites as the LiteRT export (conversion/export_vocoder.py): zero-stuffing for transposed convolutions, edge concatenation for replicate padding, one batched matrix multiply for the location-variable convolution, numeric padding instead of "same", and plain float bounds in the activation. The denoiser's STFT and the encoder's statistics are left out of the networks, because they are computed in numpy here (encode() also reads statistics with .item(), which breaks export).
  • conversion/variants.py makes the fp16 variant (ONNX Runtime's float16 converter, with float32 inputs and outputs) and the INT8 variant (ONNX Runtime dynamic quantization, QInt8 weights).

To reproduce: RESEMBLE_SRC=/path/to/resemble-enhance/src python conversion/export_onnx.py <enhancer_stage2> <out>, then python conversion/variants.py <out>.

Licence and attribution

  • Model weights and the original code: MIT, Copyright (c) 2023 Resemble AI (see LICENSE).
  • This repository: the ONNX export and the inference code are provided under the same MIT licence.
  • The mel filter bank and pipeline follow the Resemble Enhance code.
  • A LiteRT version of the same networks is in SayedShaun/resemble-enhance-litert.
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Model tree for SayedShaun/resemble-enhance-onnx

Quantized
(3)
this model