LavaSR ONNX Embedded

A PyTorch-free speech enhancement runtime for LavaSR, using ONNX Runtime with embedded model weights.

This project builds on Topping1/LavaSR-ONNX's approach of separating neural inference (ONNX) from spectral DSP (NumPy/SciPy), and adds:

• Embedded weights — model weights baked into the .onnx graph (no .data sidecar files)
• Static denoiser shapes — fixed [1, 2, 63, 257] input for 45% faster DirectML inference
• Export script — reproduce the ONNX models yourself from the upstream PyTorch checkpoint
• Benchmark results — verified performance across DirectML, CUDA, and CPU

Original LavaSR

This project is derived from ysharma3501/LavaSR. LavaSR is a lightweight speech enhancement / restoration system that supports:

• Speech bandwidth extension up to 48 kHz
• Optional denoising via ULUNAS
• Vocos-based architecture with Linkwitz-Riley spectral merge

Performance

Benchmarked on RTX 5070 Laptop GPU / Ryzen AI 9 HX 375, 2-second clip, 10 iterations after 3 warmup:

Mode	Neural Time	vs CPU
ONNX Embedded + DirectML	3.2 ms	4.8x faster
ONNX Embedded + CPU	12.0 ms	1.3x faster
PyTorch CUDA (reference)	3.8 ms	4.0x faster
PyTorch CPU (reference)	15.2 ms	baseline

Accuracy: CPU/DirectML modes produce outputs within max_diff < 0.0002 of PyTorch reference (PSNR > 102 dB).

Comparison with Topping1 Models

Model Set	DirectML	CPU	Difference
Embedded (this repo)	3.2 ms	12.0 ms	—
Topping1 (external data)	3.5 ms	12.5 ms	10% slower

The improvement comes from embedded weights (no external data file I/O) and static denoiser shapes (enables DirectML graph-level optimization).

Model Files

Three ONNX models, all with weights embedded in the graph:

File	Size	Input Shape	Description
backbone.onnx	49.4 MB	[B, 80, T] (dynamic T)	VocosBackbone: mel → hidden
spec_head.onnx	4.0 MB	[B, T, 512] (dynamic T)	ISTFTHead projection → (real, imag)
denoiser_core.onnx	1.1 MB	[1, 2, 63, 257] (static)	ULUNAS denoiser core

Only backbone.onnx and spec_head.onnx are required. The denoiser is optional (used with --denoise).

Files Needed to Run

.
├── main.py              # CLI entry point
├── lavasr_core.py       # ONNX Runtime inference + NumPy DSP
├── backbone.onnx        # VocosBackbone (from release assets)
├── spec_head.onnx       # ISTFTHead projection (from release assets)
└── denoiser_core.onnx   # ULUNAS denoiser (optional, from release assets)

Installation

Runtime Only (no PyTorch needed)

pip install numpy scipy soundfile onnxruntime

For GPU/NPU acceleration via DirectML:

pip install numpy scipy soundfile onnxruntime-directml

For Exporting Models (requires PyTorch)

pip install torch huggingface_hub LavaSR onnxruntime

Usage

Enhance Audio

# Basic enhancement
python main.py input.wav -o output.wav

# With denoising
python main.py input.wav -o output.wav --denoise

# Force DirectML provider
python main.py input.wav -o output.wav --provider directml

# Use models from a specific directory
python main.py input.wav -o output.wav --model-dir ./models

Export Models Yourself

# Export all 3 models with validation
python export_onnx.py --output-dir ./models

# Export without validation
python export_onnx.py --output-dir ./models --no-verify

Python API

from lavasr_core import LavaSROnnxRunner

runner = LavaSROnnxRunner(model_dir="./models", provider_hint="auto")
print(f"Using providers: {runner.active_providers}")

# File-based
output_path = runner.enhance_file("input.wav", output_wav="output.wav", denoise=True)

# Array-based
import soundfile as sf
audio, sr = sf.read("input.wav", dtype="float32")
enhanced, out_sr = runner.enhance_array(audio, sample_rate=sr, denoise=True)
sf.write("output.wav", enhanced, out_sr)

Architecture

The inference pipeline keeps all spectral DSP in NumPy/SciPy (no torch.stft in ONNX graphs), making the models compatible with any ONNX Runtime execution provider:

Input WAV (16 kHz)
    │
    ├─[optional]─→ STFT (NumPy) → Denoise (ONNX) → ISTFT (NumPy) → Denoised 16 kHz
    │
    ├─→ Resample to 44.1 kHz (SciPy)
    │
    ├─→ STFT + Mel filterbank (NumPy)
    │
    ├─→ Backbone (ONNX): mel → hidden
    │
    ├─→ Spec Head (ONNX): hidden → (real, imag)
    │
    ├─→ ISTFT (NumPy) → Enhanced waveform
    │
    ├─→ Spectral merge (NumPy): blend original ≤4 kHz with enhanced >4 kHz
    │
    └─→ Resample to 48 kHz (SciPy) → Output WAV

Design Decisions

Embedded Weights (vs external .data files)

By default, torch.onnx.export creates external .data files for large tensors. We use do_constant_folding=True which embeds weights directly in the ONNX protobuf. This gives ~10% DirectML speedup because the runtime doesn't need to open/mmap a separate file.

Static Denoiser Shapes (vs dynamic axes)

The denoiser is exported with fully static input shape [1, 2, 63, 257] (no dynamic axes). This allows DirectML to compile an optimized execution plan at session creation time, yielding 45% faster inference compared to dynamic shapes. The runner handles variable-length audio by chunking into 63-frame segments with overlap-add.

Two-Output Spec Head

The spec head outputs separate (real, imag) tensors instead of a single packed [B, 2050, T] tensor. This avoids complex number handling in ONNX and is compatible with all execution providers.

Acknowledgments

• ysharma3501/LavaSR — Original LavaSR model and architecture
• Topping1/LavaSR-ONNX — Pioneered the ONNX + NumPy DSP approach
• YatharthS/LavaSR — HuggingFace model weights

License

Apache-2.0 — see LICENSE.