Mazino0/conv-tasnet-onnx

Conv-TasNet Speaker Separation (ONNX)

Speaker separation model for splitting overlapping speech into individual speaker streams. Exported from SpeechBrain's pretrained SepFormer on WSJ0-2Mix.

Available Models

Model Details

Usage (Python)

import onnxruntime as ort
import numpy as np

# Use int8 model (recommended)
sess = ort.InferenceSession("conv_tasnet_libri2mix_int8.onnx")

# Input: mono waveform at 8kHz
mixture = np.random.randn(1, 8000).astype(np.float32)  # 1 second
separated = sess.run(None, {"mixture": mixture})[0]

# separated.shape = (1, 2, 8000) — two speaker sources
source_1 = separated[0, 0, :]
source_2 = separated[0, 1, :]

Usage (Rust / ort)

use ort::session::Session;
use ndarray::Array2;

let session = Session::builder()?.commit_from_file("conv_tasnet_libri2mix_int8.onnx")?;
let input = Array2::<f32>::zeros((1, 8000)); // 1s at 8kHz
let outputs = session.run(ort::inputs![input.view()])?;
// outputs[0] shape: [1, 2, 8000]

Integration in Second Brain

Used by core-asr for real-time speaker separation during detected speech overlap:

System audio (16kHz) → overlap detected?
  ├─ NO  → normal single-speaker path
  └─ YES → resample 16→8kHz → SepFormer → 2 sources → resample 8→16kHz
           → WeSpeaker embedding per source → cluster → per-speaker decode

Auto-downloaded via ensure_moonshine_models() from Mazino0/conv-tasnet-onnx.

Export & Quantization

pip install speechbrain torch onnx onnxruntime

# Export FP32 model
python scripts/export_conv_tasnet.py

# Quantize to int8
python -c "
from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic(
    'conv_tasnet_libri2mix.onnx',
    'conv_tasnet_libri2mix_int8.onnx',
    weight_type=QuantType.QUInt8,
    op_types_to_quantize=['MatMul'],
)
"

License

The SepFormer model weights are from SpeechBrain (Apache 2.0). WSJ0-2Mix data is from the Wall Street Journal corpus (LDC license).