| |
| """Export all four catalog models from checkpoints synced by sync_upstream_assets.py.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import subprocess |
| import sys |
| import tempfile |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| BASE_REVISION = "f3ff3571791e39611d31c381e3a41a3af07b4987" |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--checkpoint-manifest", |
| type=Path, |
| default=Path(tempfile.gettempdir()) / "kokoro-kmp-checkpoints.json", |
| ) |
| args = parser.parse_args() |
| if not args.checkpoint_manifest.is_file(): |
| raise SystemExit( |
| f"Missing {args.checkpoint_manifest}; run scripts/sync_upstream_assets.py first" |
| ) |
| checkpoints = json.loads(args.checkpoint_manifest.read_text(encoding="utf-8")) |
|
|
| exports = [ |
| { |
| "id": "kokoro-v1.0-fp32", |
| "checkpointName": "kokoro-v1_0.pth", |
| "config": "runtime/kokoro-v1.0-config.json", |
| "voice": "voices/en-us/af_heart.bin", |
| "output": "models/standard/kokoro-v1.0-fp32.onnx", |
| "language": "multilingual-v1.0", |
| "phonemes": "həlˈO wˈɜɹld.", |
| "repository": "hexgrad/Kokoro-82M", |
| "revision": BASE_REVISION, |
| }, |
| { |
| "id": "kokoro-v1.1-zh-fp32", |
| "checkpointName": "kokoro-v1_1-zh.pth", |
| "config": "runtime/kokoro-v1.1-zh-config.json", |
| "voice": "voices/zh/zf_001.bin", |
| "output": "models/zh/kokoro-v1.1-zh-fp32.onnx", |
| "language": "zh-en-v1.1", |
| "phonemes": "ㄋㄧ2ㄏㄠ3, ㄕ十4ㄐㄝ4.", |
| "repository": "hexgrad/Kokoro-82M-v1.1-zh", |
| "revision": "01e7505bd6a7a2ac4975463114c3a7650a9f7218", |
| }, |
| { |
| "id": "kokoro-pl-fp32", |
| "checkpointName": "finetunes/kokoro_polish_converted.pth", |
| "config": "runtime/kokoro-v1.0-config.json", |
| "voice": "voices/pl/pm_mateusz.bin", |
| "output": "models/pl/kokoro-pl-fp32.onnx", |
| "language": "pl", |
| "phonemes": "zˈaʒuwtɕ ɡˈɛɲɕlɔ̃ jˈaʒɲ.", |
| "repository": "software-mansion/react-native-executorch-kokoro", |
| "revision": "3744b57964eab7df6e8c48f0b84badb29e14df07", |
| }, |
| { |
| "id": "kokoro-de-fp32", |
| "checkpointName": "finetunes/kokoro_german_converted.pth", |
| "config": "runtime/kokoro-v1.0-config.json", |
| "voice": "voices/de/df_anna.bin", |
| "output": "models/de/kokoro-de-fp32.onnx", |
| "language": "de", |
| "phonemes": "hˈaloː vˈɛlt.", |
| "repository": "software-mansion/react-native-executorch-kokoro", |
| "revision": "3744b57964eab7df6e8c48f0b84badb29e14df07", |
| }, |
| ] |
|
|
| exporter = ROOT / "scripts/export_onnx.py" |
| for item in exports: |
| checkpoint = checkpoints.get(item["id"]) |
| if checkpoint is None: |
| raise RuntimeError(f"Checkpoint manifest has no {item['id']}") |
| command = [ |
| sys.executable, |
| str(exporter), |
| "--checkpoint", |
| checkpoint, |
| "--source-checkpoint-name", |
| item["checkpointName"], |
| "--config", |
| item["config"], |
| "--voice", |
| item["voice"], |
| "--output", |
| item["output"], |
| "--language", |
| item["language"], |
| "--reference-phonemes", |
| item["phonemes"], |
| "--source-repository", |
| item["repository"], |
| "--source-revision", |
| item["revision"], |
| "--base-revision", |
| BASE_REVISION, |
| ] |
| print(f"Exporting {item['id']}", flush=True) |
| subprocess.run(command, check=True, cwd=ROOT) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
