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
117 lines
6.1 KiB
Python
117 lines
6.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Package Plapre's AR model, Kanade codec, WavLM and HiFT in one GGUF."""
|
|
|
|
import argparse
|
|
import importlib.resources
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import torch
|
|
import torchaudio
|
|
import yaml
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def wavlm_weights(state, layers):
|
|
result = {}
|
|
for name, tensor in state.items():
|
|
if name.startswith("feature_extractor.conv_layers."):
|
|
name = name.replace(".conv.weight", ".0.weight").replace(".layer_norm.", ".2.")
|
|
elif name.startswith("encoder.feature_projection."):
|
|
name = name.replace("encoder.feature_projection.layer_norm", "layer_norm")
|
|
name = name.replace("encoder.feature_projection.projection", "post_extract_proj")
|
|
elif name.startswith("encoder.transformer.pos_conv_embed.conv."):
|
|
name = name.replace("encoder.transformer.pos_conv_embed.conv", "encoder.pos_conv.0")
|
|
name = name.replace(".parametrizations.weight.original0", ".weight_g")
|
|
name = name.replace(".parametrizations.weight.original1", ".weight_v")
|
|
elif name.startswith("encoder.transformer.layer_norm."):
|
|
name = name.replace("encoder.transformer.layer_norm", "encoder.layer_norm")
|
|
else:
|
|
match = re.fullmatch(r"encoder\.transformer\.layers\.(\d+)\.(.+)", name)
|
|
if match is None:
|
|
raise ValueError(f"Unrecognized WavLM tensor: {name}")
|
|
layer, suffix = match.groups()
|
|
if int(layer) >= layers:
|
|
continue
|
|
prefix = f"encoder.layers.{layer}."
|
|
if suffix in ("attention.attention.in_proj_weight", "attention.attention.in_proj_bias"):
|
|
parameter = "weight" if suffix.endswith("weight") else "bias"
|
|
for kind, value in zip(("q", "k", "v"), tensor.chunk(3, dim=0), strict=True):
|
|
result[prefix + f"self_attn.{kind}_proj.{parameter}"] = value.contiguous()
|
|
continue
|
|
mapping = {
|
|
"attention.gru_rel_pos_const": "self_attn.grep_a",
|
|
"attention.rel_attn_embed.weight": "self_attn.relative_attention_bias.weight",
|
|
"attention.attention.out_proj.weight": "self_attn.out_proj.weight",
|
|
"attention.attention.out_proj.bias": "self_attn.out_proj.bias",
|
|
"attention.gru_rel_pos_linear.weight": "self_attn.grep_linear.weight",
|
|
"attention.gru_rel_pos_linear.bias": "self_attn.grep_linear.bias",
|
|
"layer_norm.weight": "self_attn_layer_norm.weight",
|
|
"layer_norm.bias": "self_attn_layer_norm.bias",
|
|
"feed_forward.intermediate_dense.weight": "fc1.weight",
|
|
"feed_forward.intermediate_dense.bias": "fc1.bias",
|
|
"feed_forward.output_dense.weight": "fc2.weight",
|
|
"feed_forward.output_dense.bias": "fc2.bias",
|
|
"final_layer_norm.weight": "final_layer_norm.weight",
|
|
"final_layer_norm.bias": "final_layer_norm.bias",
|
|
}
|
|
name = prefix + mapping[suffix]
|
|
if name in result:
|
|
raise ValueError(f"Duplicate WavLM tensor: {name}")
|
|
result[name] = tensor.contiguous()
|
|
return result
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--kanade", type=Path, required=True)
|
|
parser.add_argument("--hift", type=Path, required=True)
|
|
parser.add_argument("--speakers", type=Path)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--converter", type=Path, default=Path("build/debug/bin/audiocpp_gguf"))
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
args = parser.parse_args()
|
|
source = args.source.resolve()
|
|
codec = args.kanade.resolve()
|
|
output = args.output.resolve()
|
|
spec = Path(__file__).resolve().parents[2] / "model_specs/plapre.json"
|
|
speakers = args.speakers or Path(str(importlib.resources.files("plapre") / "speakers.json"))
|
|
config = yaml.safe_load((codec / "config.yaml").read_text())["model"]["init_args"]
|
|
ssl = config["ssl_feature_extractor"]["init_args"]
|
|
if ssl["model_name"] != "wavlm_base_plus":
|
|
raise ValueError("This converter supports Kanade's WavLM Base+ frontend")
|
|
wavlm = wavlm_weights(torchaudio.pipelines.WAVLM_BASE_PLUS.get_model().state_dict(), ssl["output_layer"])
|
|
hift = torch.load(args.hift, map_location="cpu", weights_only=True)
|
|
hift = {key.removeprefix("generator."): value.contiguous() for key, value in hift.items()}
|
|
speaker_projection = torch.load(source / "speaker_proj.pt", map_location="cpu", weights_only=True)
|
|
speaker_projection = {key: value.contiguous() for key, value in speaker_projection.items()}
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix=".plapre-convert-", dir=output.parent) as temporary:
|
|
root = Path(temporary)
|
|
for name in ("config.json", "tokenizer.json", "tokenizer_config.json"):
|
|
shutil.copyfile(source / name, root / name)
|
|
shutil.copyfile(speakers, root / "speakers.json")
|
|
(root / "kanade_config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
save_file(wavlm, root / "wavlm.safetensors")
|
|
save_file(hift, root / "hift.safetensors")
|
|
save_file(speaker_projection, root / "speaker.safetensors")
|
|
command = [str(args.converter.resolve()),
|
|
"--input", str(source / "model.safetensors"),
|
|
"--input", f"speaker_proj={root / 'speaker.safetensors'}",
|
|
"--input", f"kanade={codec / 'model.safetensors'}",
|
|
"--input", f"wavlm={root / 'wavlm.safetensors'}",
|
|
"--input", f"hift={root / 'hift.safetensors'}",
|
|
"--root", str(root), "--model-spec", str(spec),
|
|
"--type", "orig", "--output", str(output)]
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
subprocess.run(command, check=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|