dfn3_mlx.py

"""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