Moipee/MossFormer2-SS-16K-CoreML

MossFormer2 speech separation (16 kHz), Core ML

A Core ML conversion of MossFormer2_SS_16K, Alibaba's speech separation model from ClearerVoice-Studio. It splits a mono 16 kHz recording of two people talking at once into one track per voice.

Inputs and outputs

Name	Type	Shape
audio	Float32, 16 kHz mono	[1, 32776] (2.05 seconds)
voices	Float32	[1, 2, 32776], one row per voice

Longer audio goes in windows of 32,776 samples that step 24,582 samples, keeping the middle of each window, as ClearerVoice-Studio does with its 2-second windows; then each voice is scaled back to the input's loudness. Which row holds which voice can change from one window to the next.

Run it on the CPU and GPU (MLComputeUnits.cpuAndGPU). The Neural Engine can't compile it, and in testing on macOS 26 one run crashed inside Core ML's CPU backend, while CPU and GPU ran without problems.

Changes from the original

• Converted to a Core ML ML program with float16 weights, for macOS 15 and later.
• A fixed input of 32,776 samples instead of any length: that's 4,096 encoder steps, 16 whole attention groups of 256, so attention needs no padding mask.
• During conversion only, two operations are written differently with the same result: padding by a negative amount becomes a crop, and einsum becomes the equivalent matrix products. PyTorch's output doesn't change.

On the same input, the Core ML model matches the PyTorch model at about 65 dB SI-SDR.

convert.py reproduces the conversion from the original weights.

License and credits

Apache 2.0, the license of the original weights and code; see LICENSE. MossFormer2 is by Shengkui Zhao and colleagues at Alibaba: MossFormer2: Combining Transformer and RNN-Free Recurrent Network for Enhanced Time-Domain Monaural Speech Separation, ICASSP 2024. All credit for the model goes to them; this repository only converts it.