79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
||
"""Prepare a Piper ONNX export for sherpa-onnx mobile inference."""
|
||
|
||
import argparse
|
||
import json
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
import onnx
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model", type=Path, required=True)
|
||
parser.add_argument("--config", type=Path, required=True)
|
||
parser.add_argument("--output-dir", type=Path, required=True)
|
||
args = parser.parse_args()
|
||
|
||
config = json.loads(args.config.read_text(encoding="utf-8"))
|
||
output_dir: Path = args.output_dir
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
output_model = output_dir / "aletheia_ru.onnx"
|
||
output_config = output_dir / "aletheia_ru.onnx.json"
|
||
output_tokens = output_dir / "tokens.txt"
|
||
|
||
shutil.copy2(args.model, output_model)
|
||
shutil.copy2(args.config, output_config)
|
||
|
||
id_map = config["phoneme_id_map"]
|
||
# sherpa-onnx's Piper lexicon maps one Unicode code point to one model ID.
|
||
# Newer Piper configs may also contain English diphthong aliases such as
|
||
# "aɪ". eSpeak emits their component code points, and sherpa rejects the
|
||
# multi-code-point aliases, so only the single-code-point table is mobile-safe.
|
||
token_rows = sorted(
|
||
((ids[0], symbol) for symbol, ids in id_map.items() if len(symbol) == 1),
|
||
key=lambda row: row[0],
|
||
)
|
||
actual_ids = [row[0] for row in token_rows]
|
||
if len(actual_ids) != len(set(actual_ids)):
|
||
raise RuntimeError("phoneme_id_map contains duplicate token identifiers")
|
||
if not actual_ids or min(actual_ids) < 0 or max(actual_ids) >= config["num_symbols"]:
|
||
raise RuntimeError("phoneme_id_map contains a token outside the model symbol range")
|
||
output_tokens.write_text(
|
||
"".join(f"{symbol} {token_id}\n" for token_id, symbol in token_rows),
|
||
encoding="utf-8",
|
||
newline="\n",
|
||
)
|
||
|
||
model = onnx.load(str(output_model))
|
||
metadata = {
|
||
"model_type": "vits",
|
||
"comment": "piper",
|
||
"language": "Russian",
|
||
"voice": config["espeak"]["voice"],
|
||
"has_espeak": "1",
|
||
"n_speakers": str(config["num_speakers"]),
|
||
"sample_rate": str(config["audio"]["sample_rate"]),
|
||
}
|
||
existing = {item.key: item for item in model.metadata_props}
|
||
for key, value in metadata.items():
|
||
if key in existing:
|
||
existing[key].value = value
|
||
else:
|
||
item = model.metadata_props.add()
|
||
item.key = key
|
||
item.value = value
|
||
onnx.save(model, str(output_model))
|
||
|
||
print(json.dumps({
|
||
"model": str(output_model),
|
||
"model_bytes": output_model.stat().st_size,
|
||
"tokens": len(token_rows),
|
||
"metadata": metadata,
|
||
}, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|