"""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] # (frames, bins) 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)