Export Kokoro 1.0 to sherpa-onnx (#1788)

scripts/kokoro/add_meta_data.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+# Copyright    2025  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import onnx
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model", type=str, required=True, help="input and output onnx model"
+    )
+
+    parser.add_argument("--voices", type=str, required=True, help="Path to voices.json")
+    return parser.parse_args()
+
+
+def load_voices(filename):
+    with open(filename) as f:
+        voices = json.load(f)
+    for key in voices:
+        voices[key] = np.array(voices[key], dtype=np.float32)
+    return voices
+
+
+def get_vocab():
+    _pad = "$"
+    _punctuation = ';:,.!?¡¿—…"«»“” '
+    _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+    dicts = {}
+    for i in range(len((symbols))):
+        dicts[symbols[i]] = i
+    return dicts
+
+
+def generate_tokens():
+    token2id = get_vocab()
+    with open("tokens.txt", "w", encoding="utf-8") as f:
+        for s, i in token2id.items():
+            f.write(f"{s} {i}\n")
+
+
+def main():
+    args = get_args()
+    print(args.model, args.voices)
+
+    model = onnx.load(args.model)
+    voices = load_voices(args.voices)
+
+    if Path("./tokens.txt").is_file():
+        print("./tokens.txt exist, skip generating it")
+    else:
+        generate_tokens()
+
+    keys = list(voices.keys())
+    print(",".join(keys))
+
+    if Path("./voices.bin").is_file():
+        print("./voices.bin exists, skip generating it")
+    else:
+        with open("voices.bin", "wb") as f:
+            for k in keys:
+                f.write(voices[k].tobytes())
+
+    speaker2id_str = ""
+    id2speaker_str = ""
+    sep = ""
+    for i, s in enumerate(keys):
+        speaker2id_str += f"{sep}{s}->{i}"
+        id2speaker_str += f"{sep}{i}->{s}"
+        sep = ","
+
+    meta_data = {
+        "model_type": "kokoro",
+        "language": "English",
+        "has_espeak": 1,
+        "sample_rate": 24000,
+        "version": 1,
+        "voice": "en-us",
+        "style_dim": ",".join(map(str, voices[keys[0]].shape)),
+        "n_speakers": len(keys),
+        "speaker2id": speaker2id_str,
+        "id2speaker": id2speaker_str,
+        "speaker_names": ",".join(keys),
+        "model_url": "https://github.com/thewh1teagle/kokoro-onnx/releases/tag/model-files",
+        "see_also": "https://huggingface.co/spaces/hexgrad/Kokoro-TTS",
+        "see_also_2": "https://huggingface.co/hexgrad/Kokoro-82M",
+        "maintainer": "k2-fsa",
+    }
+
+    print(model.metadata_props)
+
+    while len(model.metadata_props):
+        model.metadata_props.pop()
+
+    for key, value in meta_data.items():
+        meta = model.metadata_props.add()
+        meta.key = key
+        meta.value = str(value)
+    print("--------------------")
+
+    print(model.metadata_props)
+
+    onnx.save(model, args.model)
+
+    print(f"Please see {args.model}, ./voices.bin, and ./tokens.txt")
+
+
+if __name__ == "__main__":
+    main()