|
|
|
|
|
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", |
|
"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() |
|
|
