diff --git a/tests/tools/test_export_onnx.py b/tests/tools/test_export_onnx.py new file mode 100644 index 0000000..b72f34b --- /dev/null +++ b/tests/tools/test_export_onnx.py @@ -0,0 +1,66 @@ +"""Validate tools.export_onnx with a tiny synthetic model.""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from tools.export_onnx import ( + _check_flow_parity, + _check_scroll_parity, + export_flow_model, + export_scroll_decoder, +) + + +@pytest.fixture +def tiny_flow_ckpt(tmp_path: Path) -> Path: + """A flow model with seq_len=8, d_model=16, 1 layer — small but valid.""" + from tools.models import TrajectoryFlowModel + + cfg = { + "seq_len": 8, + "d_model": 16, + "nhead": 2, + "num_layers": 1, + "dim_feedforward": 32, + "cond_dim": 3, + } + model = TrajectoryFlowModel(**cfg, dropout=0.0) + model.eval() + out = tmp_path / "flow_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "flow_model.pt") + (out / "train_config.json").write_text(json.dumps(cfg)) + return out + + +@pytest.fixture +def tiny_scroll_ckpt(tmp_path: Path) -> Path: + """A scroll model with seq_len=4, latent=4, hidden=8.""" + from tools.scroll.models import ScrollCVAE + + cfg = {"seq_len": 4, "latent_dim": 4, "hidden": 8, "cond_dim": 7} + model = ScrollCVAE(**cfg) + model.eval() + out = tmp_path / "scroll_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "scroll_model.pt") + (out / "scroll_config.json").write_text(json.dumps(cfg)) + return out + + +def test_export_flow_model_parity(tiny_flow_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_flow_model(tiny_flow_ckpt, out_dir) + assert onnx_path.exists() + _check_flow_parity(tiny_flow_ckpt, onnx_path) # raises on failure + + +def test_export_scroll_decoder_parity(tiny_scroll_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_scroll_decoder(tiny_scroll_ckpt, out_dir) + assert onnx_path.exists() + _check_scroll_parity(tiny_scroll_ckpt, onnx_path)