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
85 lines
4.3 KiB
Python
85 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Package the official ClearVoice MossFormer2 separation checkpoint as GGUF."""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import gguf
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--config", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--type", choices=("f32", "f16"), default="f32")
|
|
parser.add_argument("--audiocpp-gguf", type=Path, default=ROOT / "build/debug/bin/audiocpp_gguf")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
parser.add_argument("--log", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
args.log.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(args.log), logging.StreamHandler()])
|
|
source_config = yaml.safe_load(args.config.read_text())
|
|
if source_config["network"] != "MossFormer2_SS_16K":
|
|
raise ValueError("Expected the MossFormer2 speech separation configuration")
|
|
config = {"model_type": "mossformer2"}
|
|
for key in ("sampling_rate", "num_spks", "encoder_kernel_size", "encoder_embedding_dim",
|
|
"mossformer_sequence_dim", "num_mossformer_layer"):
|
|
config[key] = source_config[key]
|
|
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
|
|
state = {key: value.clone().contiguous() for key, value in checkpoint["model"].items()}
|
|
del checkpoint
|
|
if any(value.dtype != torch.float32 for value in state.values()):
|
|
raise ValueError("Expected original F32 checkpoint tensors")
|
|
if state["enc.conv1d.weight"].shape != (
|
|
config["encoder_embedding_dim"], 1, config["encoder_kernel_size"]):
|
|
raise ValueError("Encoder configuration does not match checkpoint")
|
|
if state["mask_net.conv1d_out.weight"].shape[0] != config["num_spks"] * config["mossformer_sequence_dim"]:
|
|
raise ValueError("Speaker count does not match checkpoint")
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="mossformer2-convert-", dir=args.output.parent) as temp:
|
|
staging = Path(temp)
|
|
save_file(state, staging / "model.safetensors")
|
|
(staging / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
command = [str(args.audiocpp_gguf.resolve()), "--input", f"weights={staging / 'model.safetensors'}",
|
|
"--output", str(args.output.resolve()), "--type", "orig" if args.type == "f32" else args.type, "--family", "mossformer2",
|
|
"--model-spec", str(ROOT / "model_specs/mossformer2.json"), "--root", str(staging)]
|
|
# Affine and positional tensors feed F32 elementwise operations.
|
|
keep_f32 = {key for key, value in state.items()
|
|
if value.ndim < 2 or ".qk_offset_scale." in key}
|
|
if args.type == "f16":
|
|
for key in sorted(keep_f32):
|
|
command.extend(["--keep-type", f"weights/{key}=f32"])
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
logging.info("command=%s", command)
|
|
completed = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
|
logging.info("%s", completed.stdout)
|
|
completed.check_returncode()
|
|
reader = gguf.GGUFReader(str(args.output))
|
|
expected = {f"weights/{key}": value.numpy() for key, value in state.items()}
|
|
names = reader.fields["audiocpp.tensor_names"].contents()
|
|
if len(names) != len(reader.tensors) or set(names) != set(expected):
|
|
raise ValueError("GGUF tensor names differ from checkpoint")
|
|
for name, tensor in zip(names, reader.tensors):
|
|
original = expected[name]
|
|
if args.type == "f16" and name.removeprefix("weights/") not in keep_f32:
|
|
original = original.astype(np.float16)
|
|
if not np.array_equal(tensor.data.reshape(-1).view(np.uint8), original.reshape(-1).view(np.uint8)):
|
|
raise ValueError(f"GGUF tensor bytes differ: {name}")
|
|
logging.info("Verified %d tensors against source with %s conversion", len(expected), args.type)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|