Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import yaml
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DEFAULT_CONFIG = REPO_ROOT / "reference" / "Mel-Band-Roformer-Vocal-Model" / "configs" / "config_vocals_mel_band_roformer.yaml"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Convert Mel-Band RoFormer reference checkpoint for audio.cpp C++ runtime.")
|
|
parser.add_argument("--ckpt", default="models/melbandroformer/MelBandRoformer.ckpt")
|
|
parser.add_argument("--config-path", default=str(DEFAULT_CONFIG.relative_to(REPO_ROOT)))
|
|
parser.add_argument("--output-dir", default="models/mel-roformer-mlx")
|
|
return parser.parse_args()
|
|
|
|
|
|
def load_config(path: Path) -> dict[str, Any]:
|
|
payload = yaml.load(path.read_text(encoding="utf-8"), Loader=yaml.FullLoader)
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError(f"invalid RoFormer config payload: {path}")
|
|
return payload
|
|
|
|
|
|
def export_config_json(config: dict[str, Any]) -> dict[str, Any]:
|
|
model = dict(config["model"])
|
|
inference = dict(config["inference"])
|
|
return {
|
|
"model_type": "mel_band_roformer",
|
|
"sample_rate": int(model["sample_rate"]),
|
|
"chunk_size": int(inference["chunk_size"]),
|
|
"num_overlap": int(inference["num_overlap"]),
|
|
"dim": int(model["dim"]),
|
|
"depth": int(model["depth"]),
|
|
"num_bands": int(model["num_bands"]),
|
|
"num_stems": int(model["num_stems"]),
|
|
"time_transformer_depth": int(model["time_transformer_depth"]),
|
|
"freq_transformer_depth": int(model["freq_transformer_depth"]),
|
|
"dim_head": int(model["dim_head"]),
|
|
"heads": int(model["heads"]),
|
|
"n_fft": int(model["stft_n_fft"]),
|
|
"hop_length": int(model["stft_hop_length"]),
|
|
"win_length": int(model["stft_win_length"]),
|
|
"stft_normalized": bool(model["stft_normalized"]),
|
|
"mask_estimator_depth": int(model["mask_estimator_depth"]),
|
|
"mlp_expansion_factor": 4,
|
|
"has_final_norm": False,
|
|
}
|
|
|
|
|
|
def split_qkv_weights(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
|
converted: dict[str, torch.Tensor] = {}
|
|
for key, tensor in state.items():
|
|
if key.endswith(".rotary_embed.freqs"):
|
|
continue
|
|
if key.endswith(".to_qkv.weight"):
|
|
prefix = key[:-len(".to_qkv.weight")]
|
|
q, k, v = tensor.chunk(3, dim=0)
|
|
converted[prefix + ".to_q.weight"] = q.contiguous()
|
|
converted[prefix + ".to_k.weight"] = k.contiguous()
|
|
converted[prefix + ".to_v.weight"] = v.contiguous()
|
|
continue
|
|
converted[key] = tensor.contiguous() if isinstance(tensor, torch.Tensor) else tensor
|
|
return converted
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
ckpt_path = REPO_ROOT / args.ckpt if not Path(args.ckpt).is_absolute() else Path(args.ckpt)
|
|
config_path = REPO_ROOT / args.config_path if not Path(args.config_path).is_absolute() else Path(args.config_path)
|
|
output_dir = REPO_ROOT / args.output_dir if not Path(args.output_dir).is_absolute() else Path(args.output_dir)
|
|
|
|
state = torch.load(ckpt_path, map_location=torch.device("cpu"))
|
|
if not isinstance(state, dict):
|
|
raise RuntimeError(f"unexpected checkpoint payload type: {type(state)}")
|
|
config = load_config(config_path)
|
|
converted = split_qkv_weights(state)
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
save_file(converted, str(output_dir / "model.safetensors"))
|
|
(output_dir / "config.json").write_text(
|
|
json.dumps(export_config_json(config), indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"converted_ckpt={ckpt_path}")
|
|
print(f"output_dir={output_dir}")
|
|
print(f"tensor_count={len(converted)}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|