scripts/export_onnx.py

#!/usr/bin/env python3
"""Export a Kokoro checkpoint to a self-contained ONNX model and verify it."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch

from kokoro import KModel
from kokoro.model import KModelForONNX


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def make_reference_inputs(
    model: KModel,
    voice: Path,
    phonemes: str,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    missing = sorted({phoneme for phoneme in phonemes if model.vocab.get(phoneme) is None})
    if missing:
        raise RuntimeError(f"Reference phonemes missing from the Kokoro vocabulary: {missing}")
    if not 1 <= len(phonemes) <= 510:
        raise RuntimeError(f"Reference phoneme length must be in 1..510, got {len(phonemes)}")

    voice_pack = np.fromfile(voice, dtype=np.float32)
    if voice_pack.size != 510 * 256:
        raise RuntimeError(
            f"Expected a raw 510x256 float32 voice pack, got {voice_pack.size} values"
        )
    voice_pack = voice_pack.reshape(510, 256)

    token_ids = [int(model.vocab[phoneme]) for phoneme in phonemes]
    input_ids = torch.tensor([[0, *token_ids, 0]], dtype=torch.long)
    style = torch.from_numpy(voice_pack[len(phonemes) - 1]).unsqueeze(0)
    speed = torch.tensor([1.0], dtype=torch.float32)
    return input_ids, style, speed


def add_metadata(
    output: Path,
    *,
    language: str,
    checkpoint: Path,
    source_repository: str,
    source_checkpoint_name: str,
    source_revision: str,
    base_revision: str,
) -> None:
    model = onnx.load(output)
    onnx.helper.set_model_props(
        model,
        {
            "architecture": "Kokoro-82M",
            "language": language,
            "license": "Apache-2.0",
            "sample_rate_hz": "24000",
            "source_repository": source_repository,
            "source_revision": source_revision,
            "source_checkpoint": source_checkpoint_name,
            "source_checkpoint_sha256": sha256(checkpoint),
            "base_repository": "hexgrad/Kokoro-82M",
            "base_revision": base_revision,
        },
    )
    onnx.save_model(model, output)


def export(args: argparse.Namespace) -> None:
    checkpoint = args.checkpoint.resolve()
    config = args.config.resolve()
    voice = args.voice.resolve()
    output = args.output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(0)
    model = KModel(
        repo_id="hexgrad/Kokoro-82M",
        config=str(config),
        model=str(checkpoint),
        disable_complex=True,
    ).cpu().eval()
    export_model = KModelForONNX(model).cpu().eval()
    input_ids, style, speed = make_reference_inputs(
        model,
        voice,
        args.reference_phonemes,
    )

    torch.onnx.export(
        export_model,
        args=(input_ids, style, speed),
        f=str(output),
        input_names=["input_ids", "style", "speed"],
        output_names=["waveform", "duration"],
        opset_version=17,
        dynamic_axes={
            "input_ids": {1: "sequence_length"},
            "waveform": {0: "num_samples"},
            "duration": {0: "sequence_length"},
        },
        do_constant_folding=True,
        dynamo=False,
    )

    add_metadata(
        output,
        language=args.language,
        checkpoint=checkpoint,
        source_repository=args.source_repository,
        source_checkpoint_name=args.source_checkpoint_name or checkpoint.name,
        source_revision=args.source_revision,
        base_revision=args.base_revision,
    )
    onnx.checker.check_model(str(output), full_check=True)

    options = ort.SessionOptions()
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    session = ort.InferenceSession(str(output), sess_options=options, providers=["CPUExecutionProvider"])

    # Kokoro deliberately injects a small amount of random excitation noise. As a
    # result, PyTorch and ONNX waveforms should be highly correlated, but are not
    # expected to be bit-identical or pass a point-wise allclose assertion.
    torch.manual_seed(0)
    with torch.inference_mode():
        torch_waveform, torch_duration = export_model(input_ids, style, speed)
    ort_waveform, ort_duration = session.run(
        None,
        {
            "input_ids": input_ids.numpy(),
            "style": style.numpy(),
            "speed": speed.numpy(),
        },
    )

    expected_waveform = torch_waveform.detach().cpu().numpy()
    expected_duration = torch_duration.detach().cpu().numpy()
    np.testing.assert_array_equal(ort_duration, expected_duration)
    if not np.isfinite(ort_waveform).all():
        raise RuntimeError("ONNX Runtime produced non-finite waveform samples")

    delta = np.abs(ort_waveform - expected_waveform)
    correlation = float(np.corrcoef(ort_waveform, expected_waveform)[0, 1])
    mean_absolute_error = float(delta.mean())
    if correlation < 0.99 or mean_absolute_error > 0.005:
        raise RuntimeError(
            "ONNX/PyTorch validation failed: "
            f"correlation={correlation:.6f}, meanAbsoluteError={mean_absolute_error:.6f}"
        )
    validation = {
        "language": args.language,
        "model": output.name,
        "modelBytes": output.stat().st_size,
        "modelSha256": sha256(output),
        "sourceCheckpoint": args.source_checkpoint_name or checkpoint.name,
        "sourceCheckpointSha256": sha256(checkpoint),
        "sourceRevision": args.source_revision,
        "baseRevision": args.base_revision,
        "opset": 17,
        "sampleRateHz": 24000,
        "voice": voice.name,
        "voiceSha256": sha256(voice),
        "referencePhonemes": args.reference_phonemes,
        "referenceInputTokens": int(input_ids.shape[1]),
        "referenceOutputSamples": int(ort_waveform.size),
        "maxAbsoluteError": float(delta.max(initial=0.0)),
        "meanAbsoluteError": mean_absolute_error,
        "waveformCorrelation": correlation,
        "durationExactMatch": True,
        "onnxRuntimeVersion": ort.__version__,
        "torchVersion": torch.__version__,
    }
    validation_path = output.with_suffix(".validation.json")
    validation_path.write_text(json.dumps(validation, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(validation, indent=2))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--voice", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--language", required=True)
    parser.add_argument("--reference-phonemes", required=True)
    parser.add_argument(
        "--source-repository",
        default="software-mansion/react-native-executorch-kokoro",
    )
    parser.add_argument("--source-checkpoint-name")
    parser.add_argument("--source-revision", required=True)
    parser.add_argument("--base-revision", required=True)
    return parser.parse_args()


if __name__ == "__main__":
    export(parse_args())