Add DeepFilterNet3 MLX model weights, config, and conversion script

- model.safetensors: Pre-converted MLX weights (8.3MB)
- config.json: Model architecture configuration
- convert_deepfilternet.py: PyTorch to MLX safetensors conversion script

config.json

@@ -0,0 +1,46 @@
+{
+  "sample_rate": 48000,
+  "fft_size": 960,
+  "hop_size": 480,
+  "nb_erb": 32,
+  "nb_df": 96,
+  "df_order": 5,
+  "df_lookahead": 2,
+  "lsnr_max": 35,
+  "lsnr_min": -15,
+  "conv_ch": 64,
+  "conv_k_enc": 1,
+  "conv_k_dec": 1,
+  "conv_width_factor": 1,
+  "conv_dec_mode": "transposed",
+  "emb_hidden_dim": 256,
+  "emb_num_layers": 3,
+  "df_hidden_dim": 256,
+  "df_num_layers": 2,
+  "gru_groups": 8,
+  "linear_groups": 16,
+  "enc_linear_groups": 32,
+  "group_shuffle": false,
+  "mask_pf": false,
+  "conv_lookahead": 2,
+  "conv_depthwise": true,
+  "convt_depthwise": false,
+  "enc_concat": false,
+  "emb_gru_skip_enc": "none",
+  "emb_gru_skip": "none",
+  "df_gru_skip": "groupedlinear",
+  "dfop_method": "df",
+  "conv_kernel": [
+    1,
+    3
+  ],
+  "convt_kernel": [
+    1,
+    3
+  ],
+  "conv_kernel_inp": [
+    3,
+    3
+  ],
+  "model_version": "DeepFilterNet3"
+}

convert_deepfilternet.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+Convert DeepFilterNet PyTorch weights to MLX format.
+
+This script converts pretrained DeepFilterNet models from the original
+PyTorch implementation to MLX-compatible format with proper weight mapping.
+"""
+
+import argparse
+import json
+import configparser
+from pathlib import Path
+from typing import Dict, Any, List, Tuple
+import re
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+
+def convert_weight(weight: torch.Tensor) -> mx.array:
+    """Convert PyTorch tensor to MLX array."""
+    return mx.array(weight.detach().cpu().numpy())
+
+
+def parse_config(config_path: Path) -> Dict[str, Any]:
+    """Parse DeepFilterNet config.ini file."""
+    config = configparser.ConfigParser()
+    config.read(config_path)
+    
+    linear_groups = config.getint("deepfilternet", "linear_groups", fallback=16)
+    df_order = config.getint(
+        "df",
+        "df_order",
+        fallback=config.getint("deepfilternet", "df_order", fallback=5),
+    )
+    df_lookahead = config.getint(
+        "df",
+        "df_lookahead",
+        fallback=config.getint("deepfilternet", "df_lookahead", fallback=0),
+    )
+
+    result = {
+        # [df] section
+        "sample_rate": config.getint("df", "sr", fallback=48000),
+        "fft_size": config.getint("df", "fft_size", fallback=960),
+        "hop_size": config.getint("df", "hop_size", fallback=480),
+        "nb_erb": config.getint("df", "nb_erb", fallback=32),
+        "nb_df": config.getint("df", "nb_df", fallback=96),
+        "df_order": df_order,
+        "df_lookahead": df_lookahead,
+        "lsnr_max": config.getint("df", "lsnr_max", fallback=35),
+        "lsnr_min": config.getint("df", "lsnr_min", fallback=-15),
+        
+        # [deepfilternet] section
+        "conv_ch": config.getint("deepfilternet", "conv_ch", fallback=64),
+        "conv_k_enc": config.getint("deepfilternet", "conv_k_enc", fallback=1),
+        "conv_k_dec": config.getint("deepfilternet", "conv_k_dec", fallback=1),
+        "conv_width_factor": config.getint("deepfilternet", "conv_width_factor", fallback=1),
+        "conv_dec_mode": config.get("deepfilternet", "conv_dec_mode", fallback="transposed"),
+        "emb_hidden_dim": config.getint("deepfilternet", "emb_hidden_dim", fallback=256),
+        "emb_num_layers": config.getint("deepfilternet", "emb_num_layers", fallback=3),
+        "df_hidden_dim": config.getint("deepfilternet", "df_hidden_dim", fallback=256),
+        "df_num_layers": config.getint("deepfilternet", "df_num_layers", fallback=2),
+        "gru_groups": config.getint("deepfilternet", "gru_groups", fallback=8),
+        "linear_groups": linear_groups,
+        # DeepFilterNet2 configs do not expose enc_linear_groups separately; in that case it
+        # should follow linear_groups to keep grouped-linear tensor shapes aligned.
+        "enc_linear_groups": config.getint("deepfilternet", "enc_linear_groups", fallback=linear_groups),
+        "group_shuffle": config.getboolean("deepfilternet", "group_shuffle", fallback=False),
+        "mask_pf": config.getboolean("deepfilternet", "mask_pf", fallback=False),
+        "conv_lookahead": config.getint("deepfilternet", "conv_lookahead", fallback=2),
+        "conv_depthwise": config.getboolean("deepfilternet", "conv_depthwise", fallback=True),
+        "convt_depthwise": config.getboolean("deepfilternet", "convt_depthwise", fallback=False),
+        "enc_concat": config.getboolean("deepfilternet", "enc_concat", fallback=False),
+        "emb_gru_skip_enc": config.get("deepfilternet", "emb_gru_skip_enc", fallback="none"),
+        "emb_gru_skip": config.get("deepfilternet", "emb_gru_skip", fallback="none"),
+        "df_gru_skip": config.get("deepfilternet", "df_gru_skip", fallback="groupedlinear"),
+        "dfop_method": config.get("deepfilternet", "dfop_method", fallback="real_unfold"),
+    }
+    
+    # Parse conv_kernel strings
+    conv_kernel = config.get("deepfilternet", "conv_kernel", fallback="1,3")
+    result["conv_kernel"] = [int(x) for x in conv_kernel.split(",")]
+    
+    convt_kernel = config.get("deepfilternet", "convt_kernel", fallback="1,3")
+    result["convt_kernel"] = [int(x) for x in convt_kernel.split(",")]
+    
+    conv_kernel_inp = config.get("deepfilternet", "conv_kernel_inp", fallback="3,3")
+    result["conv_kernel_inp"] = [int(x) for x in conv_kernel_inp.split(",")]
+    
+    return result
+
+
+def convert_pytorch_to_mlx(
+    checkpoint_path: Path,
+    config_path: Path,
+    output_dir: Path,
+    model_name: str = "DeepFilterNet3",
+):
+    """Convert PyTorch checkpoint to MLX format with proper weight mapping."""
+    
+    print(f"Loading checkpoint from {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    
+    # Get state dict
+    if "state_dict" in ckpt:
+        state_dict = ckpt["state_dict"]
+    elif "model_state_dict" in ckpt:
+        state_dict = ckpt["model_state_dict"]
+    else:
+        state_dict = ckpt
+    
+    print(f"Found {len(state_dict)} parameters in checkpoint")
+    
+    # Parse config
+    print(f"Parsing config from {config_path}")
+    config_dict = parse_config(config_path)
+    config_dict["model_version"] = model_name
+    
+    # Print weight shapes for debugging
+    print("\nPyTorch weight shapes:")
+    for key, value in list(state_dict.items())[:20]:
+        print(f"  {key}: {tuple(value.shape)}")
+    print("  ...")
+    
+    # Convert weights - direct mapping since we'll match the architecture
+    print("\nConverting weights to MLX format...")
+    mlx_weights = {}
+    
+    for key, value in state_dict.items():
+        # Skip buffers that aren't needed for inference
+        if "num_batches_tracked" in key:
+            continue
+        
+        # Convert weight
+        mlx_array = convert_weight(value)
+        mlx_weights[key] = mlx_array
+    
+    print(f"Converted {len(mlx_weights)} weights")
+    
+    # Create output directory
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Save weights
+    weights_path = output_dir / "model.safetensors"
+    print(f"Saving weights to {weights_path}")
+    mx.save_safetensors(str(weights_path), mlx_weights)
+    
+    # Save config
+    config_out_path = output_dir / "config.json"
+    print(f"Saving config to {config_out_path}")
+    with open(config_out_path, "w") as f:
+        json.dump(config_dict, f, indent=2)
+    
+    print(f"\nConversion complete! Output saved to {output_dir}")
+    print(f"  - model.safetensors: {weights_path.stat().st_size / 1024 / 1024:.1f} MB")
+    print(f"  - config.json")
+    
+    return mlx_weights, config_dict
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert DeepFilterNet PyTorch weights to MLX")
+    parser.add_argument("--input", type=str, required=True, help="Path to DeepFilterNet model directory")
+    parser.add_argument("--output", type=str, required=True, help="Output directory for MLX model")
+    parser.add_argument("--name", type=str, default="DeepFilterNet3", help="Model name")
+    args = parser.parse_args()
+    
+    input_dir = Path(args.input)
+    output_dir = Path(args.output)
+    
+    # Find checkpoint
+    checkpoint_dir = input_dir / "checkpoints"
+    if checkpoint_dir.exists():
+        # Look for best checkpoint
+        checkpoints = list(checkpoint_dir.glob("*.best"))
+        if not checkpoints:
+            checkpoints = list(checkpoint_dir.glob("*.ckpt"))
+        if checkpoints:
+            checkpoint_path = checkpoints[0]
+        else:
+            raise FileNotFoundError(f"No checkpoint files found in {checkpoint_dir}")
+    else:
+        raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")
+    
+    # Find config
+    config_path = input_dir / "config.ini"
+    if not config_path.exists():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    
+    convert_pytorch_to_mlx(checkpoint_path, config_path, output_dir, args.name)
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fca0af2f25cad49d74fc9ac5f9155813416e4b350ec344bb433b4a48a9a76d38
+size 8682709