From 0ef25480cf374f864a1dce80e7a5f0c8c3c1c36f Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 00:55:15 +0800
Subject: [PATCH] feat(tools): add export_flow_model for ONNX export

---
 tools/export_onnx.py | 95 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 tools/export_onnx.py

diff --git a/tools/export_onnx.py b/tools/export_onnx.py
new file mode 100644
index 0000000..0cd4ac1
--- /dev/null
+++ b/tools/export_onnx.py
@@ -0,0 +1,95 @@
+"""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