You can deploy this model and start earning money today!

nmj21c/
MeanFlowTSE

onnx

target-speaker-extraction

Views

No views yet

ONNX Target Speaker Extraction (TSE) Models

This repository contains the ONNX-optimized versions of MeanFlowTSE (UDiT) and T-Predictor models used in the Target Speaker Extraction (TSE) application. The models are exported from PyTorch Lightning checkpoints and optimized for GPU (CUDA) and CPU deployment using ONNX Runtime.

1. Model Summary & Purpose

These models work cooperatively to extract a target speaker's voice from a noisy mixture audio:

T-Predictor (62.3 MB PyTorch -> 31.8 MB FP16 ONNX): Estimates the flow matching scaling factor ($\alpha$) by analyzing the reference target speaker voice and the mixed audio.
MeanFlowTSE (UDiT) (1.37 GB PyTorch -> 688 MB FP16 ONNX): A large Diffusion-based velocity predictor model that iteratively reconstructs clean audio features under the guide of the predicted $\alpha$.

2. File Specifications & Download List

Float32 Models (Best compatibility & fidelity):
- meanflow_tse_fp32.onnx (1375.6 MB)
- t_predictor_fp32.onnx (63.4 MB)
Float16 Models (Highly recommended for GPU):
- meanflow_tse_fp16.onnx (688.2 MB)
- t_predictor_fp16.onnx (31.8 MB)
Int8 Quantized Models (Highly optimized for CPU):
- meanflow_tse_int8.onnx (345.8 MB)
- t_predictor_int8.onnx (16.4 MB)

All evaluation-related source audio waveforms and reconstructed samples are available under the evaluation_wavs/ directory.

3. Evaluation & Performance Benchmarks

The benchmarks below were measured on an NVIDIA GeForce GPU with CUDA 12.1 and onnxruntime-gpu enabled. Numbers are averaged across 8 standard evaluation audio sets (combinations of target voices ref-voice1 to ref-voice4 and noisy mixtures mixed_noise and mixed_raw).

Model Format	Load Status	Avg Latency (ms)	Avg WAV MAE	Size Saving	Notes & Recommendations
PyTorch (Baseline)	Success	`761.1 ms`	Baseline	0% (1.43 GB)	PyTorch lightning codebase dependency
ONNX FP32	Success	`692.8 ms`	`2.71e-05`	0% (1.43 GB)	10% faster than PyTorch, mathematically equivalent
ONNX FP16	Success	`676.3 ms`	`9.91e-03`	50% saving	Best for GPU. Under 1% WAV difference, no audible noise
ONNX INT8	Success	`6716.1 ms`	`2.62e-02`	75% saving	Best for CPU (AVX512/AMX). High latency on GPU due to CPU emulations

WAV MAE: The Mean Absolute Error (MAE) calculated between PyTorch output waveforms and ONNX-reconstructed waveforms (range -1.0 to 1.0).
FP16 Quality: FP16 achieves an error rate under 1.0% compared to Float32, preserving high fidelity without introducing audible quantization noise, unlike INT8.

4. Key Fixes & Design Decisions

Timestep Embedder Type Mismatch Resolution

During direct FP16 export, PyTorch's sinusoidal embedding inside TimestepEmbedder (udit_meanflow.py) hardcoded a .float() casting, leading to a type mismatch error when multiplied by float16 weights. We modified this behavior to dynamically cast the embeddings to match the MLP's weight precision:

python

# Fixed code inside udit_meanflow.py
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
t_freq = t_freq.to(self.mlp[0].weight.dtype)  # Dynamic precision casting
t_emb = self.mlp(t_freq)

This fix enables native FP16 execution in ONNX Runtime without loading issues.

5. How to Load and Run (Python Example)

To run the models with CUDA acceleration: