File size: 7,908 Bytes
78e496a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | #!/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()
|