ai_mouse/tools/export_onnx.py

"""Export trained PyTorch checkpoints to ONNX for the inference SDK.

Usage:
    uv run python tools/export_onnx.py \
        --flow-ckpt data/models_v2 \
        --scroll-ckpt data/scroll_models \
        --output src/ai_mouse/assets/

Produces:
    <output>/flow_model.onnx
    <output>/scroll_decoder.onnx
    <output>/click_dist.json
    <output>/duration_dist.json
    <output>/train_config.json
    <output>/scroll_config.json

A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails
the .onnx files are deleted to prevent shipping broken weights.
"""
from __future__ import annotations

import argparse
import json
import logging
import shutil
import sys
from pathlib import Path

import numpy as np
import torch

logger = logging.getLogger(__name__)

_ATOL = 1e-4


def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export TrajectoryFlowModel to ONNX.

    Args:
        ckpt_dir: directory with flow_model.pt and train_config.json.
        out_dir: destination directory (created if missing).

    Returns:
        Path to the written flow_model.onnx.
    """
    from tools.models import TrajectoryFlowModel

    config_path = ckpt_dir / "train_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    d_model = int(cfg["d_model"])
    nhead = int(cfg["nhead"])
    num_layers = int(cfg["num_layers"])
    dim_feedforward = int(cfg["dim_feedforward"])
    cond_dim = int(cfg.get("cond_dim", 3))

    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        dim_feedforward=dim_feedforward,
        cond_dim=cond_dim,
        dropout=0.0,
    )
    state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "flow_model.onnx"

    dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32)
    dummy_t = torch.zeros(1, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        model,
        (dummy_x, dummy_t, dummy_cond),
        str(out_path),
        input_names=["x_t", "t", "cond"],
        output_names=["v"],
        dynamic_axes={
            "x_t": {0: "batch"},
            "t": {0: "batch"},
            "cond": {0: "batch"},
            "v": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
        dynamo=False,
    )
    logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6)
    return out_path


class _ScrollDecoder(torch.nn.Module):
    """Wraps ScrollCVAE.decode for ONNX export.

    The full ScrollCVAE is encoder+decoder; inference only needs decoder.
    """

    def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int):
        super().__init__()
        self.dec_h0 = dec_h0
        self.dec_gru = dec_gru
        self.dec_out = dec_out
        self.seq_len = seq_len
        self.hidden = hidden

    def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        b = z.shape[0]
        zc = torch.cat([z, cond], dim=-1)
        h0_flat = self.dec_h0(zc)
        h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous()
        inp = zc.unsqueeze(1).expand(b, self.seq_len, -1)
        out, _ = self.dec_gru(inp, h0)
        return self.dec_out(out)


def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export ScrollCVAE decoder to ONNX."""
    from tools.scroll.models import ScrollCVAE

    config_path = ckpt_dir / "scroll_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    latent_dim = int(cfg["latent_dim"])
    hidden = int(cfg["hidden"])
    cond_dim = int(cfg["cond_dim"])

    full = ScrollCVAE(
        seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim
    )
    state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
    full.load_state_dict(state)
    full.eval()

    decoder = _ScrollDecoder(
        dec_h0=full.dec_h0,
        dec_gru=full.dec_gru,
        dec_out=full.dec_out,
        seq_len=seq_len,
        hidden=hidden,
    )
    decoder.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "scroll_decoder.onnx"

    dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        decoder,
        (dummy_z, dummy_cond),
        str(out_path),
        input_names=["z", "cond"],
        output_names=["seq"],
        dynamic_axes={
            "z": {0: "batch"},
            "cond": {0: "batch"},
            "seq": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
        dynamo=False,
    )
    logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3)
    return out_path