File size: 1,961 Bytes
905e5e6 18a0605 905e5e6 |
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 |
#!/usr/bin/env python3
import json
import os
from typing import Any, Dict
import onnx
from iso639 import Lang
def add_meta_data(filename: str, meta_data: Dict[str, Any]):
"""Add meta data to an ONNX model. It is changed in-place.
Args:
filename:
Filename of the ONNX model to be changed.
meta_data:
Key-value pairs.
"""
model = onnx.load(filename)
for key, value in meta_data.items():
meta = model.metadata_props.add()
meta.key = key
meta.value = str(value)
onnx.save(model, filename)
def generate_tokens():
token2id = dict()
with open("phonemes.txt") as f:
for line in f:
line = line.strip()
if not line:
continue
idx, token = line.split()
if len(token) > 1:
print(f"skip token {token}")
continue
token2id[token] = int(idx)
token2id[" "] = 0
with open("tokens.txt", "w", encoding="utf-8") as f:
for s, i in token2id.items():
f.write(f"{s} {i}\n")
print("Generated tokens.txt")
def load_config(filename):
with open(filename, "r", encoding="utf-8") as file:
config = json.load(file)
return config
def main():
print("generate tokens")
generate_tokens()
lang = os.environ.get("LANG")
name = os.environ.get("NAME")
print(lang, name)
config = load_config(f"{lang}-{name}.onnx.json")
lang_iso = Lang(lang.split("_")[0])
print("add model metadata")
meta_data = {
"model_type": "vits",
"comment": "piper", # must be piper for models from piper or mimic3
"language": lang_iso.name,
"voice": lang_iso.pt1,
"has_espeak": 1,
"n_speakers": config["model"]["n_speakers"],
"sample_rate": config["audio"]["sample_rate"],
}
print(meta_data)
m = f"{lang}-{name}.onnx"
add_meta_data(m, meta_data)
main()
|