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
60 lines
3.1 KiB
Python
60 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert official FireRed VAD or Stream-VAD checkpoints to self-contained GGUF."""
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import kaldiio
|
|
import numpy as np
|
|
import torch
|
|
from safetensors.numpy import save_file
|
|
|
|
REPO = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--converter", type=Path, required=True)
|
|
parser.add_argument("--type", choices=("f32", "f16", "q8_0"), default="f32")
|
|
args = parser.parse_args()
|
|
checkpoint = torch.load(args.source / "model.pth.tar", map_location="cpu", weights_only=False)
|
|
config = vars(checkpoint["args"])
|
|
expected = dict(idim=80, odim=1, R=8, M=1, H=256, P=128, N1=20, S1=1, S2=1)
|
|
if any(config.get(key) != value for key, value in expected.items()) or config.get("N2") not in (0, 20):
|
|
raise ValueError(f"Unsupported FireRed VAD architecture: {config}")
|
|
tensors = {name: np.ascontiguousarray(value.detach().numpy(), dtype=np.float32)
|
|
for name, value in checkpoint["model_state_dict"].items()}
|
|
stats = kaldiio.load_mat(str(args.source / "cmvn.ark"))
|
|
if stats.shape != (2, 81) or stats[0, -1] < 1:
|
|
raise ValueError("Expected 80-dimensional Kaldi CMVN statistics")
|
|
mean = stats[0, :80] / stats[0, -1]
|
|
variance = np.maximum(stats[1, :80] / stats[0, -1] - mean * mean, 1e-20)
|
|
tensors["frontend.cmvn_shift"] = np.asarray(-mean, dtype=np.float32)
|
|
tensors["frontend.cmvn_scale"] = np.asarray(1 / np.sqrt(variance), dtype=np.float32)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="firered-vad-convert-") as temporary:
|
|
root = Path(temporary)
|
|
save_file(tensors, root / "weights.safetensors")
|
|
(root / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
spec = json.loads((REPO / "model_specs/firered_vad.json").read_text())
|
|
spec["modes"] = ["offline", "streaming"] if config["N2"] == 0 else ["offline"]
|
|
excluded = {"merge_silence_duration_ms", "speech_pad_ms"} if config["N2"] == 0 else {"speech_start_pad_ms"}
|
|
spec["options"]["request"] = [option for option in spec["options"]["request"] if option["name"] not in excluded]
|
|
(root / "model_spec.json").write_text(json.dumps(spec, indent=2) + "\n")
|
|
subprocess.run([str(args.converter.resolve()), "--input", f"weights={root / 'weights.safetensors'}",
|
|
"--root", str(root),
|
|
"--output", str(args.output.resolve()), "--type", args.type,
|
|
"--family", "firered_vad", "--model-spec", str(root / "model_spec.json"),
|
|
"--keep-type", "weights/frontend.*=f32",
|
|
"--overwrite"], check=True)
|
|
print(json.dumps({"variant": "Stream-VAD" if config["N2"] == 0 else "VAD",
|
|
"tensors": len(tensors), "type": args.type, "output": str(args.output)}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|