piper_train/speaker_encoder/export_encoder.py

"""Speaker Encoder ONNX Export.

Exports the ECAPA-TDNN speaker encoder from a PyTorch checkpoint to ONNX
format with dynamic batch and time axes.

Usage:
    uv run python -m piper_train.speaker_encoder.export_encoder \\
        --checkpoint speaker_encoder.ckpt \\
        --output speaker_encoder.onnx

    # With FP16 conversion (default)
    uv run python -m piper_train.speaker_encoder.export_encoder \\
        --checkpoint speaker_encoder.ckpt \\
        --output speaker_encoder.onnx \\
        --fp16

    # Without FP16 conversion
    uv run python -m piper_train.speaker_encoder.export_encoder \\
        --checkpoint speaker_encoder.ckpt \\
        --output speaker_encoder.onnx \\
        --no-fp16
"""

from __future__ import annotations

import argparse
import logging
from pathlib import Path

import torch


_LOGGER = logging.getLogger("piper_train.speaker_encoder.export_encoder")

OPSET_VERSION = 17


def export_speaker_encoder(
    checkpoint_path: Path,
    output_path: Path,
    *,
    fp16: bool = True,
    opset_version: int = OPSET_VERSION,
) -> None:
    """Export ECAPA-TDNN speaker encoder to ONNX.

    Args:
        checkpoint_path: Path to PyTorch checkpoint.
        output_path: Path for the output ONNX file.
        fp16: Whether to apply FP16 conversion (default: True).
        opset_version: ONNX opset version (default: 17).
    """
    from .ecapa_tdnn import ECAPATDNN  # noqa: PLC0415
    from .encoder import _infer_hparams  # noqa: PLC0415

    _LOGGER.info("Loading checkpoint: %s", checkpoint_path)

    # Load checkpoint
    ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=True)

    if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
        state_dict = ckpt["model_state_dict"]
    elif isinstance(ckpt, dict) and all(isinstance(k, str) for k in ckpt.keys()):
        state_dict = ckpt
    else:
        raise ValueError(
            "Checkpoint format not recognised. Expected a state_dict or a "
            "dict with 'model_state_dict' key."
        )

    hparams = _infer_hparams(state_dict)
    model = ECAPATDNN(**hparams)
    model.load_state_dict(state_dict)
    model.eval()

    _LOGGER.info(
        "Model loaded: input_dim=%d, channels=%d, emb_dim=%d",
        model.input_dim,
        model.channels,
        model.emb_dim,
    )

    # Dummy input: (batch=1, n_mels=80, time=200)
    dummy_mel = torch.randn(1, model.input_dim, 200)

    # Export
    output_path.parent.mkdir(parents=True, exist_ok=True)

    torch.onnx.export(
        model,
        (dummy_mel,),
        str(output_path),
        verbose=False,
        opset_version=opset_version,
        input_names=["mel"],
        output_names=["embedding"],
        dynamic_axes={
            "mel": {0: "batch_size", 2: "time"},
            "embedding": {0: "batch_size"},
        },
        dynamo=False,
    )

    _LOGGER.info("Exported ONNX model to %s (opset %d)", output_path, opset_version)

    # Verify the ONNX model
    _verify_onnx(output_path, model, dummy_mel)

    # FP16 conversion
    if fp16:
        _apply_fp16(output_path)

    file_size = output_path.stat().st_size
    _LOGGER.info(
        "Final model size: %.2f MB (%s)",
        file_size / (1024 * 1024),
        "FP16" if fp16 else "FP32",
    )


def _verify_onnx(
    onnx_path: Path, torch_model: torch.nn.Module, dummy_mel: torch.Tensor
) -> None:
    """Verify ONNX model produces outputs matching PyTorch.

    Args:
        onnx_path: Path to ONNX model.
        torch_model: Original PyTorch model.
        dummy_mel: Dummy input used during export.
    """
    try:
        import onnxruntime  # noqa: PLC0415
    except ImportError:
        _LOGGER.warning("onnxruntime not available; skipping ONNX verification")
        return

    import numpy as np  # noqa: PLC0415

    session = onnxruntime.InferenceSession(str(onnx_path))
    onnx_result = session.run(
        ["embedding"],
        {"mel": dummy_mel.numpy()},
    )[0]

    with torch.no_grad():
        torch_result = torch_model(dummy_mel).numpy()

    max_diff = np.abs(onnx_result - torch_result).max()
    _LOGGER.info(
        "ONNX verification: max absolute difference = %.2e (shape: %s)",
        max_diff,
        onnx_result.shape,
    )

    if max_diff > 1e-4:
        _LOGGER.warning(
            "ONNX verification: difference (%.2e) exceeds tolerance (1e-4). "
            "The exported model may have numerical issues.",
            max_diff,
        )


def _apply_fp16(onnx_path: Path) -> None:
    """Apply FP16 conversion to the ONNX model.

    Uses the project's convert_fp16 utility for VITS-compatible conversion
    that preserves numerically sensitive operators in FP32.

    Args:
        onnx_path: Path to ONNX model (modified in-place).
    """
    fp32_size = onnx_path.stat().st_size
    tmp_fp16 = onnx_path.with_suffix(".onnx.fp16_tmp")

    try:
        from ..tools.convert_fp16 import convert_fp16  # noqa: PLC0415

        convert_fp16(onnx_path, tmp_fp16)
        tmp_fp16.replace(onnx_path)
    except ImportError:
        _LOGGER.warning(
            "convert_fp16 not available; skipping FP16 conversion. "
            "The model will remain in FP32."
        )
        return
    except Exception:
        tmp_fp16.unlink(missing_ok=True)
        raise

    fp16_size = onnx_path.stat().st_size
    reduction_pct = ((fp32_size - fp16_size) / fp32_size) * 100 if fp32_size > 0 else 0
    _LOGGER.info(
        "FP16 conversion: %.2f MB -> %.2f MB (%.1f%% reduction)",
        fp32_size / (1024 * 1024),
        fp16_size / (1024 * 1024),
        reduction_pct,
    )


def main() -> None:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Export ECAPA-TDNN speaker encoder to ONNX",
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        required=True,
        help="Path to PyTorch checkpoint (.ckpt or .pt)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Path for the output ONNX model (.onnx)",
    )
    parser.add_argument(
        "--fp16",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Apply FP16 conversion (default: enabled). Use --no-fp16 to disable.",
    )
    parser.add_argument(
        "--opset",
        type=int,
        default=OPSET_VERSION,
        help=f"ONNX opset version (default: {OPSET_VERSION})",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging",
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    export_speaker_encoder(
        checkpoint_path=args.checkpoint,
        output_path=args.output,
        fp16=args.fp16,
        opset_version=args.opset,
    )


if __name__ == "__main__":
    main()