moonbox-dev/MossFormer2_SE_48K-MLX

MossFormer2 MLX

MLX builds of MossFormer2, made by Moonbox for Apple Silicon Macs. Each precision is a folder; keep a folder's files together. Only the parts listed below are converted: the rest of the pipeline comes from the original repository unless the section below says otherwise.

Files

The last column is how far each build is from the original PyTorch model on the same inputs (see the section below for what was measured).

Precision	File	Size	sha256	Checked against the original
bf16	bf16/config.json	0.4 kB	667830c83b2d	SI-SDR 11.30 dB (original 11.34), DNSMOS 3.08 (3.09); output vs the original's 45.6 dB mean / 35.5 dB worst clip
bf16	bf16/mlx_model.py	6.9 kB	65e4ee73d3bc	SI-SDR 11.30 dB (original 11.34), DNSMOS 3.08 (3.09); output vs the original's 45.6 dB mean / 35.5 dB worst clip
bf16	bf16/model.safetensors	110.7 MB	24ec32b445b3	SI-SDR 11.30 dB (original 11.34), DNSMOS 3.08 (3.09); output vs the original's 45.6 dB mean / 35.5 dB worst clip
fp16	fp16/config.json	0.4 kB	667830c83b2d	SI-SDR 11.34 dB, DNSMOS 3.09; output vs the original's 52.2 / 41.8 dB
fp16	fp16/mlx_model.py	6.9 kB	65e4ee73d3bc	SI-SDR 11.34 dB, DNSMOS 3.09; output vs the original's 52.2 / 41.8 dB
fp16	fp16/model.safetensors	110.7 MB	8b53c4ae0cba	SI-SDR 11.34 dB, DNSMOS 3.09; output vs the original's 52.2 / 41.8 dB
8bit	8bit/config.json	0.5 kB	760230e0c620	SI-SDR 11.30 dB, DNSMOS 3.08; output vs the original's 43.5 / 33.3 dB
8bit	8bit/mlx_model.py	6.9 kB	65e4ee73d3bc	SI-SDR 11.30 dB, DNSMOS 3.08; output vs the original's 43.5 / 33.3 dB
8bit	8bit/model.safetensors	67.9 MB	70ddb076cc82	SI-SDR 11.30 dB, DNSMOS 3.08; output vs the original's 43.5 / 33.3 dB
6bit	6bit/config.json	0.5 kB	c8834a4b1bed	SI-SDR 11.39 dB, DNSMOS 3.08; output vs the original's 39.9 / 26.2 dB
6bit	6bit/mlx_model.py	6.9 kB	65e4ee73d3bc	SI-SDR 11.39 dB, DNSMOS 3.08; output vs the original's 39.9 / 26.2 dB
6bit	6bit/model.safetensors	56.5 MB	6f3698c77ecb	SI-SDR 11.39 dB, DNSMOS 3.08; output vs the original's 39.9 / 26.2 dB
4bit	4bit/config.json	0.5 kB	4a2d02aa9cc4	below threshold (experimental): SI-SDR 11.43 dB mean but 0.7-0.8 dB lower on two clips, DNSMOS 3.06; output vs the original's 27.6 / 13.6 dB
4bit	4bit/mlx_model.py	6.9 kB	65e4ee73d3bc	below threshold (experimental): SI-SDR 11.43 dB mean but 0.7-0.8 dB lower on two clips, DNSMOS 3.06; output vs the original's 27.6 / 13.6 dB
4bit	4bit/model.safetensors	45.1 MB	5589ed314927	below threshold (experimental): SI-SDR 11.43 dB mean but 0.7-0.8 dB lower on two clips, DNSMOS 3.06; output vs the original's 27.6 / 13.6 dB

How to run

MossFormer2 SE 48K (ClearerVoice-Studio's speech enhancement: removes noise from speech at 48 kHz) in MLX. Each folder holds model.safetensors, config.json and mlx_model.py, the runtime: mlx-audio's MossFormer2 network (mlx-audio==0.6.5) with two corrections that make it match the original PyTorch model (below). Needs mlx, mlx-audio, numpy, soundfile (and librosa for input that isn't 48 kHz):

import sys; sys.path.insert(0, "8bit")           # a downloaded folder
import mlx_model, soundfile as sf

model = mlx_model.load("8bit")
clean = mlx_model.enhance(model, "noisy.wav")    # float32, 48 kHz; dither=0.0 makes it deterministic
sf.write("clean.wav", clean, 48000)

Why not mlx-audio's own loader. mlx-audio 0.6.5's MossFormer2SEModel loads these folders too, but its output is 18 dB (worst clip 6 dB) away from the original's and loses up to 3 dB SI-SDR: it applies rotary embeddings to the two halves of each vector (nn.RoPE(traditional=False)), where the original (rotary_embedding_torch) rotates interleaved pairs, and its Kaldi filterbank leaves the first sample of each frame un-pre-emphasized. mlx_model.py fixes both and computes the features exactly as torchaudio.compliance.kaldi.fbank (with the original's dither of 1.0 by default).

8bit/, 6bit/ and 4bit/ quantize every linear layer (MLX affine, group size 64); convolutions and norms stay bf16.

How they were checked. 8 noisy 48 kHz clips (LibriSpeech speech with babble or pink noise, 0-10 dB SNR), enhanced by the original PyTorch model (ClearerVoice's clearvoice 0.1.2, float32) and every MLX build: SI-SDR against the clean speech, DNSMOS P.835, and how close each output is to the original's (SNR in dB). Pass: SI-SDR on every clip at most 0.5 dB below the original's, mean DNSMOS at most 0.1 below.

Precision	Size	SI-SDR (dB)	DNSMOS	Output vs the original's, mean / worst clip (dB)	Real-time factor, M5 24 GB
noisy input	-	4.99	2.10	-	-
original (PyTorch fp32, MPS)	-	11.34	3.09	-	0.032
bf16/	111 MB	11.30	3.08	45.6 / 35.5	0.015
fp16/	111 MB	11.34	3.09	52.2 / 41.8	0.015
8bit/	68 MB	11.30	3.08	43.5 / 33.3	0.015
6bit/	57 MB	11.39	3.08	39.9 / 26.2	0.015
4bit/ (experimental)	45 MB	11.43 (0.7-0.8 lower on two clips)	3.06	27.6 / 13.6	0.015

With the dither off on both sides, the MLX float32 network matches the original to 57 dB (48 dB worst clip).

How these were made

From MossFormer2 at revision eff8c97925: the original weights loaded in MLX and stored as bf16 or fp16, or quantized with MLX's affine quantization (8bit, 6bit, 4bit, group size 64; layers MLX can't quantize, such as convolutions and norms, stay bf16). Every build was checked against the original before upload; the sha256 column above is what the files must hash to.

License

Apache-2.0. These files are a derivative of the original model and carry its license: it is copied in this repository, and it governs your use of these files. The original model card follows unchanged.

---

Original model card

Introduction

The MossFormer2_SE_48K model weights for 48 kHz speech enhancement in ClearerVoice-Studio repo.

This model is trained on large scale datasets inclduing open-sourced and private data.

It enhances speech audios by removing background noise.

Install

Clone the Repository

git clone https://github.com/modelscope/ClearerVoice-Studio.git

Create Conda Environment

cd ClearerVoice-Studio
conda create -n clearvoice python=3.8
conda activate clearvoice
pip install -r requirements.txt

Run Script

Go to clearvoice/ and use the following examples. The MossFormer2_SE_48K model will be downloaded from huggingface automatically.

Sample example 1: use speech enhancement model MossFormer2_SE_48K to process one wave file of samples/input.wav and save the output wave file to samples/output_MossFormer2_SE_48K.wav

from clearvoice import ClearVoice

myClearVoice = ClearVoice(task='speech_enhancement', model_names=['MossFormer2_SE_48K'])

output_wav = myClearVoice(input_path='samples/input.wav', online_write=False)

myClearVoice.write(output_wav, output_path='samples/output_MossFormer2_SE_48K.wav')

Sample example 2: use speech enhancement model MossFormer2_SE_48K to process all input wave files in samples/path_to_input_wavs/ and save all output files to samples/path_to_output_wavs

from clearvoice import ClearVoice

myClearVoice = ClearVoice(task='speech_enhancement', model_names=['MossFormer2_SE_48K'])

myClearVoice(input_path='samples/path_to_input_wavs', online_write=True, output_path='samples/path_to_output_wavs')

Sample example 3: use speech enhancement model MossFormer2_SE_48K to process wave files listed in `samples/audio_samples.scp' file, and save all output files to 'samples/path_to_output_wavs_scp/'

from clearvoice import ClearVoice

myClearVoice = ClearVoice(task='speech_enhancement', model_names=['MossFormer2_SE_48K'])

myClearVoice(input_path='samples/scp/audio_samples.scp', online_write=True, output_path='samples/path_to_output_wavs_scp')