File size: 2,611 Bytes
7a83907 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
#!/usr/bin/env python3
import json
import os
from typing import Any, Dict
import onnx
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 load_config(model):
with open(f"{model}.json", "r") as file:
config = json.load(file)
return config
def generate_tokens(config):
id_map = config["phoneme_id_map"]
with open("tokens.txt", "w", encoding="utf-8") as f:
for s, i in id_map.items():
f.write(f"{s} {i[0]}\n")
print("Generated tokens.txt")
lang_map = {
"da_DK": "Danish",
"el_GR": "Greek",
"fi_FI": "Finnish",
"hu_HU": "Hungarian",
"is_IS": "Icelandic",
"it_IT": "Italian",
"ka_GE": "Georgian",
"kk_KZ": "Kazakh",
"lb_LU": "Luxembourgish",
"ne_NP": "Nepali",
"nl_BE": "Dutch",
"nl_NL": "Dutch",
"no_NO": "Norwegian",
"pl_PL": "Polish",
"pt_BR": "Portuguese",
"pt_PT": "Portuguese",
"ro_RO": "Romanian",
"sk_SK": "Slovak",
"sr_RS": "Serbian",
"sv_SE": "Swedish",
"sw_CD": "Swahili",
"tr_TR": "Turkish",
"vi_VN": "Vietnamese",
}
# for en_US-lessac-medium.onnx
# export LANG=en_US
# export TYPE=lessac
# export NAME=medium
def main():
lang = os.environ.get("LANG", None)
if not lang:
print("Please provide the environment variable LANG")
return
t = os.environ.get("TYPE", None)
if not t:
print("Please provide the environment variable TYPE")
return
# thorsten or thorsten_emotional
name = os.environ.get("NAME", None)
if not t:
print("Please provide the environment variable NAME")
return
print("type", t)
config = load_config(f"{lang}-{name}-{t}.onnx")
print("generate tokens")
generate_tokens(config)
print("add model metadata")
meta_data = {
"model_type": "vits",
"comment": "piper", # must be piper for models from piper
"language": lang_map[lang],
"voice": config["espeak"]["voice"], # e.g., en-us
"has_espeak": 1,
"n_speakers": config["num_speakers"],
"sample_rate": config["audio"]["sample_rate"],
}
print(meta_data)
add_meta_data(f"{lang}-{name}-{t}.onnx", meta_data)
main()
|