"""Pure-MLX reference implementation of the DeepFilterNet3 network. Loads the model.safetensors + config.json shipped in this repository and mirrors the DeepFilterNet3 inference graph: lookahead shift on the input features, encoder, ERB decoder, DF decoder, DF-output reshape. The surrounding DSP (STFT, ERB feature extraction, normalization, deep-filter application, iSTFT) is the caller's job — auxiliary.npz carries the exact filterbank / window / normalization constants. Layout: channels-last [B, T, F, C] throughout (time = conv H axis, frequency = conv W axis). All GRUs run with a zero initial hidden state. """ import json from pathlib import Path import mlx.core as mx import numpy as np def _relu(x): return mx.maximum(x, 0) def grouped_linear(x, w): """x: [B, T, I], w: [groups, I/groups, H/groups] -> [B, T, H].""" b, t, _ = x.shape g, gi, gh = w.shape x = x.reshape(b, t, g, gi) y = mx.einsum("btgi,gih->btgh", x, w) return y.reshape(b, t, g * gh) def gru_layer(x, Wx, Wh, b, bhn): """mlx.nn.GRU recurrence with an explicit zero initial hidden state. x: [B, T, I] -> [B, T, H]. Gate order r, z, n (same as PyTorch): r = sigmoid(x@Wxr.T + h@Whr.T + br) z = sigmoid(x@Wxz.T + h@Whz.T + bz) n = tanh(x@Wxn.T + bn + r * (h@Whn.T + bhn)) h' = (1 - z) * n + z * h """ hidden = Wh.shape[1] xp = x @ Wx.T + b # [B, T, 3H] x_rz = xp[..., : 2 * hidden] x_n = xp[..., 2 * hidden :] h = mx.zeros((x.shape[0], hidden), dtype=x.dtype) outs = [] for t in range(x.shape[1]): hp = h @ Wh.T rz = mx.sigmoid(x_rz[:, t] + hp[..., : 2 * hidden]) r = rz[..., :hidden] z = rz[..., hidden:] n = mx.tanh(x_n[:, t] + r * (hp[..., 2 * hidden :] + bhn)) h = (1 - z) * n + z * h outs.append(h) return mx.stack(outs, axis=1) def conv_transpose_dw_fstride2(x, w_flipped): """Depthwise ConvTranspose2d over frequency, k=(1,3), stride=(1,2), padding=(0,1), output_padding=(0,1) — via zero-stuffing + correlation. x: [B, T, F, C] -> [B, T, 2F, C]. w_flipped: [C, 1, 3, 1], kernel pre-flipped along the frequency axis. """ bsz, t, f, c = x.shape z = mx.stack([x, mx.zeros_like(x)], axis=3).reshape(bsz, t, 2 * f, c) z = mx.pad(z, [(0, 0), (0, 0), (1, 1), (0, 0)]) return mx.conv2d(z, w_flipped, stride=(1, 1), padding=(0, 0), groups=c) class DFN3MLX: """DeepFilterNet3 network. Inputs are normalized features: feat_erb [B, T, 32, 1] — ERB features (dB-scaled, mean-normalized) feat_spec [B, T, 96, 2] — complex spectrogram features (unit-normalized) Returns (erb_mask [B, T, 32, 1], df_coefs [B, 5, T, 96, 2], lsnr [B, T, 1]). """ def __init__(self, model_dir: str | Path): model_dir = Path(model_dir) self.w = {k: mx.array(v) for k, v in mx.load(str(model_dir / "model.safetensors")).items()} with open(model_dir / "config.json") as f: self.cfg = json.load(f) # Pre-flip transposed-conv kernels for the zero-stuff formulation for name in ("erb_dec.convt2.dwt.weight", "erb_dec.convt1.dwt.weight"): flipped = np.array(self.w[name])[:, :, ::-1, :] self.w[name] = mx.array(np.ascontiguousarray(flipped)) self.lsnr_scale = self.cfg["lsnr_max"] - self.cfg["lsnr_min"] self.lsnr_offset = self.cfg["lsnr_min"] self.conv_lookahead = self.cfg["conv_lookahead"] # ── building blocks ── def _inp_conv(self, x, prefix, groups): """pad(t=2) + Conv2d k=(3,3) (grouped) [+ pointwise] + ReLU.""" w = self.w x = mx.pad(x, [(0, 0), (2, 0), (0, 0), (0, 0)]) if f"{prefix}.conv.bias" in w: # erb_conv0: single conv, BN fused into it y = mx.conv2d(x, w[f"{prefix}.conv.weight"], stride=(1, 1), padding=(0, 1), groups=groups) y = y + w[f"{prefix}.conv.bias"] else: # df_conv0: grouped conv + BN-fused pointwise y = mx.conv2d(x, w[f"{prefix}.conv.weight"], stride=(1, 1), padding=(0, 1), groups=groups) y = mx.conv2d(y, w[f"{prefix}.pw.weight"], stride=(1, 1), padding=(0, 0)) y = y + w[f"{prefix}.pw.bias"] return _relu(y) def _sep_conv(self, x, prefix, fstride): """Depthwise k=(1,3) conv + BN-fused pointwise + ReLU.""" w = self.w c = x.shape[-1] y = mx.conv2d(x, w[f"{prefix}.dw.weight"], stride=(1, fstride), padding=(0, 1), groups=c) y = mx.conv2d(y, w[f"{prefix}.pw.weight"], stride=(1, 1), padding=(0, 0)) return _relu(y + w[f"{prefix}.pw.bias"]) def _pathway(self, x, prefix): """Depthwise 1x1 conv (BN fused) + ReLU — a per-channel affine.""" w = self.w y = mx.conv2d(x, w[f"{prefix}.dw.weight"], stride=(1, 1), padding=(0, 0), groups=x.shape[-1]) return _relu(y + w[f"{prefix}.dw.bias"]) def _sep_convt(self, x, prefix): """Depthwise transposed conv (fstride=2) + BN-fused pointwise + ReLU.""" w = self.w y = conv_transpose_dw_fstride2(x, w[f"{prefix}.dwt.weight"]) y = mx.conv2d(y, w[f"{prefix}.pw.weight"], stride=(1, 1), padding=(0, 0)) return _relu(y + w[f"{prefix}.pw.bias"]) def _squeezed_gru(self, x, prefix, num_layers, has_linear_out): w = self.w y = _relu(grouped_linear(x, w[f"{prefix}.linear_in.weight"])) if num_layers == 1: y = gru_layer(y, w[f"{prefix}.gru.Wx"], w[f"{prefix}.gru.Wh"], w[f"{prefix}.gru.b"], w[f"{prefix}.gru.bhn"]) else: for l in range(num_layers): p = f"{prefix}.gru.layers.{l}" y = gru_layer(y, w[f"{p}.Wx"], w[f"{p}.Wh"], w[f"{p}.b"], w[f"{p}.bhn"]) if has_linear_out: y = _relu(grouped_linear(y, w[f"{prefix}.linear_out.weight"])) return y # ── model ── def __call__(self, feat_erb, feat_spec, apply_lookahead_shift: bool = True): w = self.w la = self.conv_lookahead if apply_lookahead_shift and la > 0: feat_erb = mx.pad(feat_erb[:, la:], [(0, 0), (0, la), (0, 0), (0, 0)]) feat_spec = mx.pad(feat_spec[:, la:], [(0, 0), (0, la), (0, 0), (0, 0)]) # Encoder e0 = self._inp_conv(feat_erb, "enc.erb_conv0", groups=1) # [B,T,32,64] e1 = self._sep_conv(e0, "enc.erb_conv1", 2) # [B,T,16,64] e2 = self._sep_conv(e1, "enc.erb_conv2", 2) # [B,T,8,64] e3 = self._sep_conv(e2, "enc.erb_conv3", 1) # [B,T,8,64] c0 = self._inp_conv(feat_spec, "enc.df_conv0", groups=2) # [B,T,96,64] c1 = self._sep_conv(c0, "enc.df_conv1", 2) # [B,T,48,64] bsz, t = e3.shape[0], e3.shape[1] cemb = _relu(grouped_linear(c1.reshape(bsz, t, -1), w["enc.df_fc_emb.weight"])) emb = e3.reshape(bsz, t, -1) + cemb # [B,T,512] emb = self._squeezed_gru(emb, "enc.emb_gru", num_layers=1, has_linear_out=True) lsnr = mx.sigmoid(emb @ w["enc.lsnr_fc.weight"].T + w["enc.lsnr_fc.bias"]) lsnr = lsnr * self.lsnr_scale + self.lsnr_offset # [B,T,1] # ERB decoder d = self._squeezed_gru(emb, "erb_dec.emb_gru", num_layers=2, has_linear_out=True) d = d.reshape(bsz, t, 8, -1) # [B,T,8,64] d3 = self._sep_conv(self._pathway(e3, "erb_dec.conv3p") + d, "erb_dec.convt3", 1) d2 = self._sep_convt(self._pathway(e2, "erb_dec.conv2p") + d3, "erb_dec.convt2") d1 = self._sep_convt(self._pathway(e1, "erb_dec.conv1p") + d2, "erb_dec.convt1") m = self._pathway(e0, "erb_dec.conv0p") + d1 m = mx.conv2d(m, w["erb_dec.conv0_out.conv.weight"], stride=(1, 1), padding=(0, 1)) erb_mask = mx.sigmoid(m + w["erb_dec.conv0_out.conv.bias"]) # [B,T,32,1] # DF decoder c = self._squeezed_gru(emb, "df_dec.df_gru", num_layers=2, has_linear_out=False) c = c + grouped_linear(emb, w["df_dec.df_skip.weight"]) # [B,T,256] cp = mx.pad(c0, [(0, 0), (4, 0), (0, 0), (0, 0)]) cp = mx.conv2d(cp, w["df_dec.df_convp.conv.weight"], stride=(1, 1), padding=(0, 0), groups=2) cp = mx.conv2d(cp, w["df_dec.df_convp.pw.weight"], stride=(1, 1), padding=(0, 0)) cp = _relu(cp + w["df_dec.df_convp.pw.bias"]) # [B,T,96,10] cf = mx.tanh(grouped_linear(c, w["df_dec.df_out.weight"])) # [B,T,960] cf = cf.reshape(bsz, t, 96, 10) + cp df_coefs = cf.reshape(bsz, t, 96, 5, 2).transpose(0, 3, 1, 2, 4) # [B,5,T,96,2] return erb_mask, df_coefs, lsnr