File size: 1,609 Bytes
0c31dc5 |
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 |
#!/usr/bin/env python3
import json
import os
import re
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
# for en_US-lessac-medium.onnx
# export TYPE=lessac
# export NAME=medium
def main():
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"en_US-{name}-{t}.onnx")
print("add model metadata")
meta_data = {
"model_type": "vits",
"comment": "piper", # must be piper for models from piper
"language": "English",
"voice": config["espeak"]["voice"], # e.g., en-us
"has_espeak": int("espeak" in config),
"n_speakers": config["num_speakers"],
"sample_rate": config["audio"]["sample_rate"],
}
print(meta_data)
add_meta_data(f"en_US-{name}-{t}.onnx", meta_data)
main()
|