| """MossFormer2 SE 48K in MLX: speech enhancement (denoising) at 48 kHz, as ClearerVoice-Studio runs it. |
| |
| The network is mlx-audio's port (mlx_audio.sts.models.mossformer2_se, mlx-audio 0.6.5) with two corrections that |
| make it match the original PyTorch model: |
| |
| - rotary embeddings on interleaved pairs (rotary_embedding_torch, which the original uses); mlx-audio builds |
| nn.RoPE(traditional=False), which rotates the two halves instead |
| - Kaldi filterbanks computed exactly as torchaudio.compliance.kaldi.fbank: pre-emphasis of each frame's first |
| sample against itself (x0 - 0.97 x0), log floor at float32 epsilon, dither as an explicit argument (the original |
| uses dither=1.0; 0 makes the output deterministic) |
| |
| import mlx_model |
| model = mlx_model.load("8bit") # a folder: config.json + model.safetensors |
| clean = mlx_model.enhance(model, "noisy.wav") # float32 numpy, 48 kHz (input resampled to 48 kHz) |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
| import numpy as np |
| from mlx.utils import tree_unflatten |
|
|
| MAX_WAV_VALUE = 32768.0 |
| SAMPLE_RATE = 48000 |
|
|
|
|
| def load(folder) -> "nn.Module": |
| from mlx_audio.sts.models.mossformer2_se.config import MossFormer2SEConfig |
| from mlx_audio.sts.models.mossformer2_se.mossformer2_se_wrapper import MossFormer2SE |
| folder = Path(folder) |
| conf = json.loads((folder / "config.json").read_text()) |
| q = conf.pop("quantization_config", None) |
| conf.pop("model_type", None) |
| cfg = MossFormer2SEConfig.from_dict(conf) |
| model = MossFormer2SE(cfg) |
| for _, m in model.named_modules(): |
| if isinstance(getattr(m, "rotary_pos_emb", None), nn.RoPE): |
| r = m.rotary_pos_emb |
| m.rotary_pos_emb = nn.RoPE(r.dims, traditional=True, base=r.base) |
| if q: |
| nn.quantize(model, group_size=q["group_size"], bits=q["bits"]) |
| model.update(tree_unflatten(list(mx.load(str(folder / "model.safetensors")).items()))) |
| mx.eval(model.parameters()) |
| model.eval() |
| model.config = cfg |
| return model |
|
|
|
|
| def _hamming(n: int) -> mx.array: |
| k = mx.arange(n, dtype=mx.float32) |
| return 0.54 - 0.46 * mx.cos(2 * np.pi * k / (n - 1)) |
|
|
|
|
| def fbank(x: mx.array, cfg, dither: float = 1.0) -> mx.array: |
| """torchaudio.compliance.kaldi.fbank(x, dither, frame_length=win_len, frame_shift=win_inc, num_mel_bins=num_mels, |
| sample_frequency=48000, window_type='hamming') with Kaldi's other defaults (snip_edges, remove_dc_offset, |
| preemphasis 0.97, power spectrum, low_freq 20, high_freq Nyquist).""" |
| from mlx_audio.dsp import get_mel_banks_kaldi |
| size, shift = cfg.win_len, cfg.win_inc |
| n = 1 + (x.shape[0] - size) // shift |
| idx = mx.arange(n)[:, None] * shift + mx.arange(size)[None, :] |
| frames = x[idx] |
| if dither: |
| frames = frames + mx.random.normal(frames.shape) * dither |
| frames = frames - mx.mean(frames, axis=1, keepdims=True) |
| prev = mx.concatenate([frames[:, :1], frames[:, :-1]], axis=1) |
| frames = (frames - 0.97 * prev) * _hamming(size) |
| pad = 1 << (size - 1).bit_length() |
| spec = mx.abs(mx.fft.rfft(frames, n=pad, axis=1)) ** 2 |
| banks, _ = get_mel_banks_kaldi(cfg.num_mels, pad, float(SAMPLE_RATE), 20.0, 0.0) |
| banks = mx.pad(banks, [(0, 0), (0, 1)]) |
| return mx.log(mx.maximum(spec @ banks.T, np.finfo(np.float32).eps)) |
|
|
|
|
| def _deltas(x: mx.array) -> mx.array: |
| """torchaudio.functional.compute_deltas(win_length=5) over the last axis (replicate padding).""" |
| p = mx.concatenate([x[..., :1], x[..., :1], x, x[..., -1:], x[..., -1:]], axis=-1) |
| t = x.shape[-1] |
| return (-2 * p[..., 0:t] - p[..., 1:t + 1] + p[..., 3:t + 3] + 2 * p[..., 4:t + 4]) / 10.0 |
|
|
|
|
| def _stft(x: mx.array, cfg) -> mx.array: |
| """The original's stft (torch.stft, center=False, hamming window of win_len, n_fft = fft_len): (frames, bins).""" |
| size, shift = cfg.win_len, cfg.win_inc |
| n = 1 + (x.shape[0] - size) // shift |
| idx = mx.arange(n)[:, None] * shift + mx.arange(size)[None, :] |
| return mx.fft.rfft(x[idx] * _hamming(size), n=cfg.fft_len, axis=1) |
|
|
|
|
| def _istft(spec: np.ndarray, cfg, length: int) -> np.ndarray: |
| """torch.istft(center=False, hamming window): overlap-add with window-square normalisation.""" |
| size, shift = cfg.win_len, cfg.win_inc |
| w = np.array(_hamming(size)) |
| frames = np.fft.irfft(spec, n=cfg.fft_len, axis=1)[:, :size] * w |
| n = frames.shape[0] |
| out = np.zeros(size + shift * (n - 1)); norm = np.zeros_like(out) |
| for i in range(n): |
| out[i * shift:i * shift + size] += frames[i] |
| norm[i * shift:i * shift + size] += w ** 2 |
| nz = norm > 1e-11 |
| out[nz] /= norm[nz] |
| return np.pad(out, (0, max(0, length - len(out))))[:length] |
|
|
|
|
| def _segment(model, audio: np.ndarray, dither: float) -> np.ndarray: |
| cfg = model.config |
| x = mx.array(audio.astype(np.float32)) |
| f = fbank(x, cfg, dither) |
| d1 = _deltas(f.T); d2 = _deltas(d1) |
| feats = mx.concatenate([f, d1.T, d2.T], axis=1)[None] |
| mask = model(feats)[-1][0] |
| spec = _stft(x, cfg) |
| n = min(spec.shape[0], mask.shape[0]) |
| masked = np.array(spec[:n] * mask[:n].astype(mx.float32)) |
| return _istft(masked, cfg, len(audio)) |
|
|
|
|
| def enhance(model, audio, sr: int | None = None, dither: float = 1.0) -> np.ndarray: |
| """Denoise `audio` (a path or a float waveform at `sr`); returns float32 at 48 kHz, the input's length. Inputs over |
| 20 s go through the original's 4 s windows with 3 s stride (the edges of each window dropped).""" |
| if isinstance(audio, (str, Path)): |
| import soundfile as sf |
| audio, sr = sf.read(str(audio), dtype="float32") |
| audio = np.asarray(audio, dtype=np.float32) |
| audio = audio.mean(1) if audio.ndim > 1 else audio |
| if sr and sr != SAMPLE_RATE: |
| import librosa |
| audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE) |
| cfg = model.config |
| length = len(audio) |
| x = audio * MAX_WAV_VALUE |
| if length <= SAMPLE_RATE * cfg.one_time_decode_length: |
| return (_segment(model, x, dither) / MAX_WAV_VALUE).astype(np.float32) |
| window = int(SAMPLE_RATE * cfg.decode_window) |
| stride = int(window * 0.75) |
| t = len(x) |
| if t < window: |
| x = np.concatenate([x, np.zeros(window - t)]) |
| elif t < window + stride: |
| x = np.concatenate([x, np.zeros(window + stride - t)]) |
| elif (t - window) % stride: |
| x = np.concatenate([x, np.zeros(t - (t - window) // stride * stride)]) |
| out = np.zeros(len(x)) |
| give = (window - stride) // 2 |
| i = 0 |
| while i + window <= len(x): |
| seg = _segment(model, x[i:i + window], dither) |
| if i == 0: |
| out[:window - give] = seg[:-give] |
| else: |
| out[i + give:i + window - give] = seg[give:-give] |
| i += stride |
| return (out[:length] / MAX_WAV_VALUE).astype(np.float32) |
|
|
