| --- |
| 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. |
|
|
