Add conversion script and update README with conversion instructions

README.md

@@ -53,15 +53,16 @@ All versions share the same audio parameters:
 ## Files
 
 ```
+convert_deepfilternet.py   # PyTorch → MLX conversion script
 v1/
-  config.json          # v1 architecture configuration
-  model.safetensors    # v1 weights
+  config.json              # v1 architecture configuration
+  model.safetensors        # v1 weights
 v2/
-  config.json          # v2 architecture configuration
-  model.safetensors    # v2 weights
+  config.json              # v2 architecture configuration
+  model.safetensors        # v2 weights
 v3/
-  config.json          # v3 architecture configuration
-  model.safetensors    # v3 weights
+  config.json              # v3 architecture configuration
+  model.safetensors        # v3 weights
 ```
 
 ## Usage
@@ -90,11 +91,29 @@ let model = try await DeepFilterNetModel.fromPretrained("mlx-community/DeepFilte
 let enhanced = try model.enhance(audioArray)
 ```
 
+## Converting from PyTorch
+
+To re-create these weights from the original DeepFilterNet checkpoints:
+
+```bash
+# Clone the original repo to get the pretrained checkpoints
+git clone https://github.com/Rikorose/DeepFilterNet
+
+# Convert each version
+python convert_deepfilternet.py --input DeepFilterNet/DeepFilterNet  --output v1 --name DeepFilterNet
+python convert_deepfilternet.py --input DeepFilterNet/DeepFilterNet2 --output v2 --name DeepFilterNet2
+python convert_deepfilternet.py --input DeepFilterNet/DeepFilterNet3 --output v3 --name DeepFilterNet3
+```
+
+Each input directory should contain a `config.ini` and a `checkpoints/` folder from the original repo.
+
+Requires `torch` and `mlx` to be installed.
+
 ## Origin
 
 - **Original model:** [DeepFilterNet](https://github.com/Rikorose/DeepFilterNet) by Hendrik Schroeter
 - **License:** MIT (same as the original)
-- **Conversion:** PyTorch → `safetensors`
+- **Conversion:** PyTorch → `safetensors` via `convert_deepfilternet.py`
 
 ## Citations
 

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()