---
license: apache-2.0
tags:
- target-speaker-extraction
- speech-enhancement
- audio-processing
- onnx
- onnxruntime
- flow-matching
- udit
- t-predictor
---

# 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:

1. **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.
2. **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:

```python
import onnxruntime as ort
import numpy as np

# Select CUDA Execution Provider
providers = [('CUDAExecutionProvider', {'device_id': 0}), 'CPUExecutionProvider']

# Load FP16 sessions
tp_session = ort.InferenceSession("t_predictor_fp16.onnx", providers=providers)
udit_session = ort.InferenceSession("meanflow_tse_fp16.onnx", providers=providers)

# Example: T-Predictor Inference
# Inputs: mixture (Batch, Time_Steps), enrollment (Batch, Time_Steps)
mixture_data = np.random.randn(1, 48000).astype(np.float16)  # Use float16 for FP16 models
enroll_data = np.random.randn(1, 48000).astype(np.float16)

tp_inputs = {"mixture": mixture_data, "enrollment": enroll_data}
tp_outputs = tp_session.run(None, tp_inputs)
alpha = tp_outputs[0]
print("Predicted Alpha:", alpha)
```

---

## 6. Original Sources & Checkpoints
* **Original Codebase**: MeanFlowTSE uses Diffusion-based Velocity Predictor architectures derived from DiT (Diffusion Transformers) and GLIDE concepts.
  * UDiT Module reference: [GLIDE Text2Im](https://github.com/openai/glide-text2im)
* **Checkpoint Origin**: The weights exported into these ONNX models were trained on clean-speech TSE datasets and extracted from the following local training checkpoints:
  * MeanFlowTSE (UDiT) baseline: `backend/exp/best-clean-weights.ckpt`
  * T-Predictor baseline: `backend/exp/t-predictor-clean-weights.ckpt`

---

## 7. License
This repository is licensed under the Apache License 2.0. Feel free to use, modify, and distribute these models in your projects. Refer to the `LICENSE` file for details.