#!/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()