Merge branch 'worktree-lib-refactor': ai_mouse 0.2.0 library refactor

Pure-inference ONNX Runtime SDK; src/ layout; bundled weights;
training/server/eval moved to tools/; new public API
(MouseModel/ScrollModel + cached generate/generate_scroll helpers).
See CHANGELOG.md for details.

.github/workflows/ci.yml

@@ -0,0 +1,37 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  library:
+    name: Library tests (no torch)
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        python: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+      - run: uv venv --python ${{ matrix.python }}
+      - run: uv pip install -e . pytest
+      - run: uv run pytest tests/unit -v
+
+  dev:
+    name: Full dev suite (with torch)
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        python: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+      - run: uv sync --group dev --python ${{ matrix.python }}
+      - run: uv run pytest tests/ -v

CHANGELOG.md

@@ -0,0 +1,39 @@
+# Changelog
+
+All notable changes to this project will be documented here. Format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows
+[Semantic Versioning](https://semver.org/).
+
+## [0.2.0] - 2026-05-12
+
+### Changed (breaking)
+
+- Inference no longer requires PyTorch. Runtime dependencies are now
+  `numpy + onnxruntime` only.
+- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a
+  persistent ORT `InferenceSession`.
+- Function signatures `generate()` and `generate_scroll()` are now keyword-only
+  past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`).
+- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll).
+- Removed `config=` parameter; use kwargs directly.
+- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`.
+- `start_scrollY` / `target_scrollY` renamed to `start_scroll_y` / `target_scroll_y`.
+- Training, web UI, collector, eval, and data adapter code moved to repo-level
+  `tools/`; no longer packaged in the wheel.
+
+### Added
+
+- ONNX-format pre-trained weights bundled inside the wheel via
+  `importlib.resources` (~3 MB).
+- `tools/export_onnx.py` script with PyTorch/ORT parity check.
+- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`,
+  `GenerationError`.
+- Custom ORT providers parameter for GPU / DirectML.
+- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the
+  default model across calls.
+
+### Removed
+
+- Legacy `JointCVAE` class.
+- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs).
+- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling).

CLAUDE.md

@@ -4,63 +4,103 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 ## Project
 
-Local FastAPI tool that trains and serves ML models for generating **human-like mouse trajectories and scroll wheel events**. Frontend is Vue 3 + ECharts loaded from CDN. Package manager is **uv** (Python 3.12–3.13).
+`ai_mouse` is an ONNX-Runtime SDK that generates human-like mouse trajectories
+and scroll wheel events. Runtime dependencies are `numpy + onnxruntime` only;
+training and the FastAPI web UI live under `tools/` and are not packaged.
+
+Package manager: **uv**, Python 3.12-3.13.
+
+## Library vs tools — hard boundary
+
+- **`src/ai_mouse/`** — wheel content. NEVER add `import torch` /
+  `import fastapi` / `import scipy` / `import matplotlib` here. CI's
+  `library` job installs only runtime deps and would break.
+- **`tools/`** — repo-only dev code (training, server, collector, eval,
+  data adapters, ONNX export). May `import` library private modules
+  (`ai_mouse._coord`, `ai_mouse._postprocess`) freely — they co-evolve.
+- **Bundled assets**: `src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx`
+  plus four JSON metadata files. Re-generated by
+  `tools/export_onnx.py` after retraining.
 
 ## Commands
 
 ```bash
-# Run the web app (opens http://127.0.0.1:8765 in browser)
+# Web UI (collect + train + verify in browser)
-uv run python main.py
+uv run python tools/serve.py
 
-# Tests (httpx + pytest-asyncio for ASGI integration tests)
+# Tools CLI dispatch
-uv run pytest
+uv run python -m tools train --data data/traces.jsonl --output data/models_v2
-uv run pytest tests/test_generator.py
+uv run python -m tools eval --model-dir data/models_v2 \
-uv run pytest tests/test_server.py::TestStatus::test_status_returns_trace_count
+    --reference data/pretrain_traces.jsonl --output data/eval_reports/r.md
+uv run python -m tools balabit-adapter --input data/balabit_raw \
+    --output data/pretrain_traces.jsonl
 
-# Sync dependencies
+# Re-export ONNX (after retraining)
-uv sync
+uv run python -m tools.export_onnx --flow-ckpt data/models_v2 \
+    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
+
+# Tests
+uv run pytest tests/unit            # library-only (no torch)
+uv run pytest tests/tools           # full dev suite (needs [dev] group)
+uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility
+
+# Dependency sync
+uv sync                              # runtime only
+uv sync --group dev                  # dev-everything
 ```
 
-There is no separate lint/format config; do not invent one.
-
 ## Architecture
 
-Two parallel ML subsystems share an identical **collect → train → verify** workflow. Both pipelines persist JSONL traces, train via SSE-streamed progress, and load bundled weights for inference.
+Two parallel ML subsystems share a **collect → train → export → serve** flow.
 
-### Mouse trajectories (`ai_mouse/`)
+### Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training)
 
-- **Model**: `TrajectoryFlowModel` — Conditional Flow Matching (OT) with a 4-layer pre-norm Transformer backbone. Trained by sampling `t ~ U[0,1]`, interpolating `x_t = (1-t)·noise + t·data`, and regressing the velocity field `v = x1 - x0` via MSE.
+- **Model**: `TrajectoryFlowModel` (Conditional Flow Matching with 4-layer
-- **Inference** ([generator.py](ai_mouse/generator.py)): 10-step Euler ODE from noise → trajectory in rotated frame → heavy spatial/temporal post-processing (endpoint snapping, forward monotonicity, log_dt clipping, asymmetric speed profile) → pixel decoding → cumulative timestamps → appended click-down/click-up events sampled from a truncated normal.
+  pre-norm Transformer, d_model=128, defined in `tools/models.py`)
-- **Rotated coordinate frame** ([coord.py](ai_mouse/coord.py)): All trajectories are normalised so `start → (0, 0)` and `end → (1, 0)`. Lateral is perpendicular. This makes the model **angle- and distance-invariant** — a 1000px diagonal and a 200px horizontal look identical to the network. `encode_trajectory` / `decode_trajectory` are the only bridge between pixel space and model space.
+- **Inference**: 10-step Euler ODE in Python; each step runs
-- **Condition vector** (3 dims): `[dist/2000, log(dist/100), log(total_dur/500)]`. `total_duration` is sampled at inference from a per-distance-bin log-normal stored in `duration_dist.json`.
+  `session.run(...)` on `src/ai_mouse/assets/flow_model.onnx`. Followed by
-- **Training artefacts** in `data/models_v2/`: `flow_model.pt`, `click_dist.json` (truncated-normal click duration), `duration_dist.json` (per-bin log-normal), `train_config.json` (architecture params — required for inference to instantiate the model with matching hyperparameters).
+  numpy post-processing in `_postprocess.py` (endpoint snapping, forward
-- **6× data augmentation** in [trainer.py](ai_mouse/trainer.py): original, lateral flip, ±20% speed, temporal noise, flip+speed.
+  monotonicity, gaussian smoothing, log_dt → cumulative timestamps,
-- **Legacy** `JointCVAE` in `models.py` is kept only for backward compatibility; the active model is `TrajectoryFlowModel`.
+  truncated-normal click duration).
+- **Rotated coordinate frame** (`_coord.py`): trajectories normalised so
+  `start → (0, 0)`, `end → (1, 0)`. Makes the model angle/distance invariant.
 
-### Scroll wheel (`ai_mouse/scroll/`)
+### Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`)
 
-- **Model**: `ScrollCVAE` — smaller bidirectional-GRU encoder + GRU decoder VAE. Sequences are 32 wheel events of `(delta_norm, log_Δt)`.
+- **Model**: `ScrollCVAE` (bidirectional-GRU encoder + GRU decoder VAE,
-- **Condition vector** (7 dims): `[dist/5000, log(dist/500), direction, viewport_norm, mode_onehot×3]` where mode is `"target" | "fast" | "precise"` (see `SCROLL_MODES` in [config.py](ai_mouse/config.py)).
+  `tools/scroll/models.py`). Only the **decoder** is exported to ONNX
-- **Inference**: VAE prior sample → softmax-normalise deltas to sum to ~1 → scale to target distance → quantise to wheel increments (40px precise, 120px otherwise) → adjust last event so total matches exactly.
+  (`scroll_decoder.onnx`); encoder is training-only.
-- **Note**: scroll *collection* state is JS-side (the browser fires real `wheel` events); the Python `ScrollCollector` only generates targets and persists traces.
+- **Inference**: sample `z ~ N(0, 1)` in numpy → one `session.run(...)` →
+  softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) →
+  scale to target distance → cumulative timestamps.
 
-### Server (`ai_mouse/server/`)
+### Server (`tools/server/`) and frontend (`static/`)
 
-- App factory `create_app()` in [server/__init__.py](ai_mouse/server/__init__.py) mounts four routers under `/api`: `routes_collect`, `routes_train`, `routes_verify`, and `routes_scroll` (the scroll router uses prefix `/api/scroll`).
+App factory `create_app()` mounts four routers under `/api`. Frontend is
-- **Session state** ([server/deps.py](ai_mouse/server/deps.py)): A module-level `SessionState` singleton holds the active `Collector` / `ScrollCollector`. Tests that need to override `_DATA_DIR` monkeypatch `ai_mouse.server.deps._DATA_DIR`.
+vanilla Vue 3 + axios + ECharts via CDN. Note: the `/api/verify` and
-- **Training progress** is delivered via Server-Sent Events. The route launches `train()` on a background thread (`asyncio.to_thread`) and pushes `{epoch, total, loss}` dicts through an `asyncio.Queue` until `{done: True}` or `{error: ...}`.
+`/api/scroll/verify` endpoints always use the **bundled** ONNX weights
-- **Static** is mounted from `static/` at the project root (not under the package); `index.html` is served at `/`.
+(via `from ai_mouse import generate / generate_scroll`). If you retrain
-
+and want the Web UI to reflect new weights, re-run `tools.export_onnx` and
-### Frontend (`static/`)
+restart the server.
-
-Vanilla Vue 3 + axios + ECharts pulled from CDN — **no bundler, no node_modules**. Single page with three tab views (`collect.js`, `train.js`, `verify.js`) registered as components in [app.js](static/js/app.js). `api.js` exports an axios instance and a `fetchSSE()` helper that reads `text/event-stream` from `fetch().body.getReader()` and parses `data: ...` frames. UI strings are in Chinese.
 
 ## Config
 
-[ai_mouse/config.py](ai_mouse/config.py) is the **single source of truth** for hyperparameters. `TrainConfig`, `GenerateConfig`, `ScrollTrainConfig`, `ScrollModeConfig`, and `ServerConfig` are dataclasses. When changing model architecture (`d_model`, `nhead`, etc.) keep training and inference consistent — the `train_config.json` saved at training time is what `generate()` uses to reconstruct the model.
+`tools/config.py` holds the training-side dataclasses (`TrainConfig`,
+`ScrollTrainConfig`, etc.). The library does NOT use these — its only
+"configuration" is what's embedded in `src/ai_mouse/assets/train_config.json`
+(architecture params needed to know `seq_len` etc. at inference time).
 
 ## Tests
 
-`tests/conftest.py` provides `model_dir` and `scroll_model_dir` fixtures that write **freshly initialised (untrained)** weights to a temp dir along with all required JSON metadata. Use these whenever a test calls `generate()` or `generate_scroll()`. The trajectories will be garbage but the inference path runs end-to-end.
+- `tests/unit/conftest.py` — fixtures for library-only tests, no torch.
+- `tests/tools/conftest.py` — `model_dir` and `scroll_model_dir` fixtures
+  that produce **untrained** torch weights in a temp dir. Used by training-
+  /server-side tests.
+- `tests/unit/test_golden.py` — regression suite that pins library output
+  against `tests/unit/data/golden_{mouse,scroll}.npz` captured before the
+  ONNX migration. Tolerance is distance-scaled: mouse allows ±max(30 px, 20%
+  of move distance) and ±700 ms; scroll requires exact total deltaY match
+  and ±2 quanta per event.
 
-Server tests use `httpx.ASGITransport(app=create_app())` with `pytest-asyncio` — no live socket.
+Server tests use `httpx.ASGITransport(app=create_app())` with
+`pytest-asyncio` — no live socket.

README.md

@@ -0,0 +1,120 @@
+# ai_mouse
+
+Human-like mouse trajectory and scroll wheel event generator. Inference runs on
+ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`.
+
+## Install
+
+```bash
+pip install git+https://github.com/<owner>/ai_mouse.git
+```
+
+For GPU inference (optional), replace `onnxruntime` with the GPU variant:
+
+```bash
+pip install onnxruntime-gpu      # CUDA / TensorRT
+# or
+pip install onnxruntime-directml # Windows DirectML
+```
+
+## Quick start
+
+### Mouse trajectory
+
+```python
+from ai_mouse import generate
+
+points = generate(start=(100, 200), end=(900, 400))
+# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)]
+```
+
+### Scroll wheel
+
+```python
+from ai_mouse import generate_scroll
+
+events = generate_scroll(start_scroll_y=0, target_scroll_y=2000)
+# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...]
+```
+
+### Class API (recommended for repeated calls)
+
+```python
+from ai_mouse import MouseModel
+
+m = MouseModel()              # session created once
+for target in target_list:
+    pts = m.generate((cx, cy), target)
+```
+
+### Custom providers / GPU
+
+```python
+from ai_mouse import MouseModel
+
+m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+# or
+m = MouseModel(providers=["DmlExecutionProvider"])
+```
+
+### Reproducibility
+
+```python
+m.generate(start, end, seed=42)
+```
+
+## API summary
+
+| Name | Purpose |
+|---|---|
+| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton |
+| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session |
+| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll |
+| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy |
+
+## Thread safety
+
+`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently
+from multiple threads — ORT `InferenceSession` is itself thread-safe.
+
+## Development
+
+The repo contains optional dev-only tooling under `tools/` for training your
+own models, running the FastAPI web UI, and evaluating output quality. Install
+with the `dev` group:
+
+```bash
+uv sync --group dev
+```
+
+Common commands:
+
+```bash
+# Web UI (collect + train + verify in browser)
+uv run python tools/serve.py
+
+# Training (after collecting your own data)
+uv run python -m tools train --data data/traces.jsonl --output data/models_v2
+
+# Convert Balabit corpus to trace format
+uv run python -m tools balabit-adapter --input data/balabit_raw \
+    --output data/pretrain_traces.jsonl
+
+# Eval report
+uv run python -m tools eval --model-dir data/models_v2 \
+    --reference data/pretrain_traces.jsonl --output data/eval_reports/report.md
+
+# Re-export ONNX after retraining
+uv run python -m tools.export_onnx --flow-ckpt data/models_v2 \
+    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
+```
+
+Tests:
+
+```bash
+uv run pytest tests/unit    # library-only (no torch)
+uv run pytest tests/tools   # full dev suite
+```
+
+After retraining you need to re-export and rebuild the wheel for the new
+weights to ship; the in-app Verify endpoint always uses bundled weights.

ai_mouse/__init__.py

@@ -1,5 +0,0 @@
-# sites/ai_mouse/ai_mouse/__init__.py
-from ai_mouse.generator import generate
-from ai_mouse.scroll.generator import generate_scroll
-
-__all__ = ["generate", "generate_scroll"]

ai_mouse/generator.py

@@ -1,353 +0,0 @@
-"""Inference layer: Flow Matching trajectory generation.
-
-Pipeline:
-1. Load model from model_dir (flow_model.pt, click_dist.json,
-   duration_dist.json, train_config.json).
-2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)].
-3. Sample total_duration from duration_dist.json by distance bin (log-normal).
-4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory.
-5. Spatial post-processing:
-   a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
-   b. Smooth start: dampen lateral near start (first 4 points).
-   c. Enforce forward monotonicity (prevent x-axis jitter).
-   d. 5-point gaussian smooth on lateral (preserve endpoints).
-6. Temporal post-processing:
-   a. Clip log_dt to [0, 5] to prevent exponential explosion.
-   (speed profile and median±1.1 hard clip removed in 2026-05 refactor —
-    let the model's learned timing distribution come through naturally.)
-7. Decode to pixels via decode_trajectory.
-8. Resample to n_points if n_points != model seq_len.
-9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150].
-10. Ensure timestamps monotonically increasing.
-11. Append click events sampled from truncated normal.
-"""
-from __future__ import annotations
-
-import json
-import logging
-import math
-from pathlib import Path
-
-import numpy as np
-import torch
-from scipy.stats import truncnorm
-
-from ai_mouse.config import GenerateConfig
-from ai_mouse.coord import decode_trajectory
-from ai_mouse.models import TrajectoryFlowModel
-from ai_mouse.utils import resample_arc
-
-logger = logging.getLogger(__name__)
-
-_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / "models_v2"
-
-
-# ---------------------------------------------------------------------------
-# Duration sampling helper
-# ---------------------------------------------------------------------------
-
-
-def _sample_duration(duration_dist: dict, dist: float) -> float:
-    """Sample a total movement duration (ms) for the given pixel distance.
-
-    Uses per-distance-bin log-normal parameters from duration_dist.
-
-    Args:
-        duration_dist: dict with "bins" and "params" keys.
-        dist: pixel distance between start and end.
-
-    Returns:
-        Sampled duration in milliseconds.
-    """
-    bins = duration_dist["bins"]
-    params = duration_dist["params"]
-    # Find bin for this distance
-    bin_idx = len(bins) - 1
-    for i in range(len(bins) - 1):
-        if dist < bins[i + 1]:
-            bin_idx = i
-            break
-    # Clamp to valid params index
-    bin_idx = min(bin_idx, len(params) - 1)
-    mu_log = params[bin_idx]["mu_log"]
-    sigma_log = params[bin_idx]["sigma_log"]
-    return float(np.exp(np.random.normal(mu_log, sigma_log)))
-
-
-# ---------------------------------------------------------------------------
-# Smoothing helper
-# ---------------------------------------------------------------------------
-
-
-def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
-    """5-point gaussian smoothing along a 1-D array, preserving endpoints.
-
-    Args:
-        x: 1-D input array.
-        sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
-               weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].
-
-    Returns:
-        Smoothed array of the same shape. x[0] and x[-1] are unchanged.
-        If len(x) < 5, returns x unchanged (kernel won't fit).
-    """
-    if len(x) < 5:
-        return x.copy()
-    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
-    kernel /= kernel.sum()
-    # Pad with edge values to avoid boundary artifacts, then slice back
-    padded = np.pad(x, pad_width=2, mode="edge")
-    smoothed = np.convolve(padded, kernel, mode="valid")
-    smoothed[0] = x[0]
-    smoothed[-1] = x[-1]
-    return smoothed
-
-
-# ---------------------------------------------------------------------------
-# Main generate function
-# ---------------------------------------------------------------------------
-
-
-def generate(
-    start: tuple[int, int],
-    end: tuple[int, int],
-    n_points: int = 64,
-    speed: float | None = None,
-    model_dir: str | None = None,
-    config: GenerateConfig | None = None,
-) -> list[tuple[int, int, int]]:
-    """Generate a human-like mouse trajectory from start to end.
-
-    Uses a Flow Matching model with 4-step Euler ODE integration.
-
-    Args:
-        start:      (x, y) starting pixel coordinate.
-        end:        (x, y) target pixel coordinate.
-        n_points:   number of movement points in the path (default 64).
-        speed:      optional speed multiplier; speed=2 halves the duration.
-        model_dir:  directory containing flow_model.pt, click_dist.json,
-                    duration_dist.json, train_config.json.
-                    None → use bundled pre-trained weights.
-        config:     GenerateConfig instance; None → use defaults.
-
-    Returns:
-        List of (x, y, t_ms) tuples.  All values are ints.
-        Last two entries are the mouse-down and mouse-up click events.
-    """
-    if config is None:
-        config = GenerateConfig()
-
-    model_dir_path = Path(model_dir) if model_dir else _BUNDLED_MODELS_DIR
-
-    flow_pt = model_dir_path / "flow_model.pt"
-    click_json = model_dir_path / "click_dist.json"
-    duration_json = model_dir_path / "duration_dist.json"
-    config_json = model_dir_path / "train_config.json"
-
-    if not flow_pt.exists():
-        if model_dir is not None:
-            raise FileNotFoundError(
-                f"Model weights not found in {model_dir_path}. "
-                "Run training first or omit model_dir to use bundled weights."
-            )
-        raise FileNotFoundError(
-            f"Bundled model weights missing at {_BUNDLED_MODELS_DIR}. "
-            "Run training first."
-        )
-
-    # Load train config for model architecture params
-    seq_len = config.seq_len
-    d_model = 128
-    nhead = 4
-    num_layers = 4
-    dim_feedforward = 256
-    cond_dim = 3
-    if config_json.exists():
-        cfg = json.loads(config_json.read_text())
-        seq_len = int(cfg.get("seq_len", seq_len))
-        d_model = int(cfg.get("d_model", d_model))
-        nhead = int(cfg.get("nhead", nhead))
-        num_layers = int(cfg.get("num_layers", num_layers))
-        dim_feedforward = int(cfg.get("dim_feedforward", dim_feedforward))
-        cond_dim = int(cfg.get("cond_dim", cond_dim))
-
-    # Load model
-    model = TrajectoryFlowModel(
-        seq_len=seq_len,
-        d_model=d_model,
-        nhead=nhead,
-        num_layers=num_layers,
-        dim_feedforward=dim_feedforward,
-        cond_dim=cond_dim,
-    )
-    model.load_state_dict(
-        torch.load(flow_pt, map_location="cpu", weights_only=True)
-    )
-    model.eval()
-
-    # Load auxiliary distributions
-    click_params: dict = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
-    if click_json.exists():
-        click_params = json.loads(click_json.read_text())
-
-    duration_dist: dict | None = None
-    if duration_json.exists():
-        duration_dist = json.loads(duration_json.read_text())
-
-    # Compute pixel distance
-    sx, sy = float(start[0]), float(start[1])
-    ex, ey = float(end[0]), float(end[1])
-    dist = math.hypot(ex - sx, ey - sy)
-    dist = max(dist, 1.0)
-
-    # Sample total duration
-    if duration_dist is not None:
-        total_duration = _sample_duration(duration_dist, dist)
-    else:
-        # Fallback: simple heuristic ~2px/ms
-        total_duration = dist / 2.0
-    if speed is not None and speed > 0:
-        total_duration /= speed
-    total_duration = max(total_duration, 10.0)
-
-    # Build condition vector: [dist_norm, log_dist, log_total_dur]
-    cond_arr = np.array(
-        [
-            dist / 2000.0,
-            math.log(dist / 100.0),
-            math.log(total_duration / 500.0),
-        ],
-        dtype=np.float32,
-    )
-    cond_t = torch.from_numpy(cond_arr).unsqueeze(0)  # (1, 3)
-
-    # -----------------------------------------------------------------------
-    # 4-step Euler ODE integration
-    # -----------------------------------------------------------------------
-    n_steps = config.n_steps
-    dt = 1.0 / n_steps
-
-    with torch.no_grad():
-        x = torch.randn(1, seq_len, 3)  # start from noise
-        for step in range(n_steps):
-            t_val = step * dt
-            t_tensor = torch.tensor([t_val])
-            v = model(x, t_tensor, cond_t)
-            x = x + v * dt
-        # x is now the generated trajectory in (forward, lateral, log_dt) space
-
-    decoded = x.squeeze(0).numpy()  # (seq_len, 3)
-
-    forward = decoded[:, 0].copy()  # (seq_len,)
-    lateral = decoded[:, 1].copy()  # (seq_len,)
-    log_dt = decoded[:, 2].copy()   # (seq_len,)
-
-    # ------------------------------------------------------------------
-    # Spatial post-processing
-    # ------------------------------------------------------------------
-
-    # Endpoint snapping: lerp last 6 points towards (1.0, 0.0)
-    n_snap = min(6, seq_len // 4)
-    for i in range(n_snap):
-        alpha = ((i + 1) / n_snap) ** 2  # quadratic ease-in
-        k = seq_len - n_snap + i
-        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
-        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
-
-    # Force first and last points to canonical values
-    forward[0], lateral[0] = 0.0, 0.0
-    forward[-1], lateral[-1] = 1.0, 0.0
-
-    # Smooth start: dampen lateral near start (first 4 points)
-    n_start_fix = min(4, seq_len // 4)
-    for i in range(1, n_start_fix + 1):
-        blend = i / (n_start_fix + 1)  # 0.2, 0.4, 0.6, 0.8
-        forward[i] = max(forward[i], forward[i - 1])  # ensure monotonic start
-        lateral[i] = lateral[i] * blend  # dampen lateral near start
-
-    # Enforce forward monotonicity with soft correction (prevent x-jitter)
-    for i in range(1, seq_len - 1):  # skip last point (already snapped to 1.0)
-        if forward[i] < forward[i - 1]:
-            forward[i] = forward[i - 1] + 0.001
-
-    # Clamp forward to [0, 1] and re-force endpoints after monotonicity fix
-    forward = np.clip(forward, 0.0, 1.0)
-    forward[0] = 0.0
-    forward[-1] = 1.0
-
-    # Lateral 5-point gaussian smoothing (endpoints preserved)
-    lateral = _gaussian_smooth(lateral, sigma=1.0)
-
-    # ------------------------------------------------------------------
-    # Temporal post-processing (log_dt)
-    # ------------------------------------------------------------------
-
-    # Clip log_dt to prevent extreme values after exp()
-    # Training data log_dt = log(Δt_ms + 1), typical range [0, 4.5]
-    # (e.g., Δt=1ms → 0.69, Δt=10ms → 2.40, Δt=80ms → 4.39)
-    log_dt = np.clip(log_dt, 0.0, 5.0)
-
-    # First point has no interval (padding from training)
-    log_dt[0] = 0.0
-
-    # Decode spatial coordinates to pixels
-    normalised = np.stack([forward, lateral], axis=1)  # (seq_len, 2)
-    pixels = decode_trajectory(normalised, start, end)  # (seq_len, 2)
-
-    # Resample to n_points if needed
-    if n_points != seq_len:
-        pixels = resample_arc(pixels, n_points)
-        # Also resample log_dt via linear interpolation in uniform arc
-        log_dt = np.interp(
-            np.linspace(0, 1, n_points),
-            np.linspace(0, 1, seq_len),
-            log_dt,
-        )
-
-    xs = pixels[:, 0]
-    ys = pixels[:, 1]
-
-    # Convert log_dt → dt (ms), scale to total_duration, clip
-    dt_raw = np.exp(log_dt)
-    dt_raw = np.clip(dt_raw, 0.0, None)
-    dt_sum = dt_raw.sum()
-    if dt_sum > 1e-6:
-        scale = total_duration / dt_sum
-    else:
-        scale = total_duration / max(n_points, 1)
-    dt_ms = np.clip(
-        dt_raw * scale,
-        config.dt_clip_min_ms,
-        config.dt_clip_max_ms,
-    )
-
-    # Cumulative timestamps (start at 0)
-    t_abs = np.cumsum(dt_ms)
-    t_abs = np.concatenate([[0.0], t_abs[:-1]])  # shift so first point = 0
-
-    # Ensure monotonically increasing
-    for i in range(1, len(t_abs)):
-        if t_abs[i] <= t_abs[i - 1]:
-            t_abs[i] = t_abs[i - 1] + 1.0
-
-    move_points: list[tuple[int, int, int]] = [
-        (int(round(xs[i])), int(round(ys[i])), int(round(t_abs[i])))
-        for i in range(n_points)
-    ]
-
-    # Sample click duration from truncated normal
-    mu = float(click_params["mu"])
-    sigma = float(click_params["sigma"])
-    low = float(click_params["low"])
-    high = float(click_params["high"])
-    a, b = (low - mu) / sigma, (high - mu) / sigma
-    click_duration = int(truncnorm.rvs(a, b, loc=mu, scale=sigma))
-    click_duration = max(click_duration, int(low))
-
-    last_t = move_points[-1][2]
-    click_x = int(round(xs[-1]))
-    click_y = int(round(ys[-1]))
-    return move_points + [
-        (click_x, click_y, last_t),
-        (click_x, click_y, last_t + click_duration),
-    ]

ai_mouse/scroll/__init__.py

@@ -1,6 +0,0 @@
-"""Scroll wheel event generation subsystem."""
-from ai_mouse.scroll.generator import generate_scroll
-from ai_mouse.scroll.trainer import train_scroll
-from ai_mouse.scroll.collector import ScrollCollector
-
-__all__ = ["generate_scroll", "train_scroll", "ScrollCollector"]

ai_mouse/scroll/generator.py

@@ -1,148 +0,0 @@
-"""Scroll wheel event sequence generator."""
-from __future__ import annotations
-
-import json
-import logging
-import math
-from pathlib import Path
-
-import numpy as np
-import torch
-
-from ai_mouse.scroll.models import ScrollCVAE
-
-logger = logging.getLogger(__name__)
-
-_BUNDLED_SCROLL_MODELS = Path(__file__).resolve().parent.parent.parent / "data" / "scroll_models"
-
-
-def _build_condition(
-    distance: float,
-    direction: int,
-    mode: str,
-    viewport_height: float = 900.0,
-) -> np.ndarray:
-    """Build 7-dim condition vector matching the trainer layout.
-
-    Dims: [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3]
-    """
-    mode_onehot = [0.0, 0.0, 0.0]
-    if mode == "target":
-        mode_onehot[0] = 1.0
-    elif mode == "fast":
-        mode_onehot[1] = 1.0
-    elif mode == "precise":
-        mode_onehot[2] = 1.0
-
-    viewport_norm = viewport_height / 1000.0
-
-    return np.array([
-        distance / 5000.0,
-        math.log(max(distance, 1.0) / 500.0),
-        float(direction),
-        viewport_norm,
-        *mode_onehot,
-    ], dtype=np.float32)
-
-
-def generate_scroll(
-    start_scrollY: int,
-    target_scrollY: int,
-    mode: str = "target",
-    model_dir: str | None = None,
-) -> list[dict]:
-    """Generate a realistic scroll event sequence.
-
-    Args:
-        start_scrollY: Current scroll position (px from top).
-        target_scrollY: Target scroll position.
-        mode: "target" | "fast" | "precise"
-        model_dir: Path to scroll model files. None = bundled.
-
-    Returns:
-        List of {"deltaY": int, "deltaMode": 0, "t": int}.
-        Positive deltaY = scroll down, negative = scroll up.
-    """
-    model_dir_path = Path(model_dir) if model_dir else _BUNDLED_SCROLL_MODELS
-    model_pt = model_dir_path / "scroll_model.pt"
-    config_json = model_dir_path / "scroll_config.json"
-
-    if not model_pt.exists():
-        raise FileNotFoundError(f"Scroll model not found at {model_pt}")
-
-    seq_len = 32
-    if config_json.exists():
-        cfg = json.loads(config_json.read_text())
-        seq_len = cfg.get("seq_len", 32)
-
-    model = ScrollCVAE(seq_len=seq_len)
-    model.load_state_dict(torch.load(model_pt, map_location="cpu", weights_only=True))
-    model.eval()
-
-    distance = abs(target_scrollY - start_scrollY)
-    direction = 1 if target_scrollY > start_scrollY else -1
-    distance = max(distance, 10)
-
-    cond = _build_condition(float(distance), direction, mode)
-    cond_t = torch.from_numpy(cond).unsqueeze(0)
-
-    with torch.no_grad():
-        z = torch.randn(1, model.latent_dim)
-        decoded = model.decode(z, cond_t).squeeze(0).numpy()
-
-    delta_norm = decoded[:, 0]
-    log_dt = decoded[:, 1]
-
-    # De-normalise delta: use softmax-like normalisation so they sum to ~1
-    delta_weights = np.exp(delta_norm)
-    delta_weights = delta_weights / delta_weights.sum()
-    delta_px = delta_weights * distance * direction
-
-    # Quantise to realistic wheel increments
-    quantum = 40 if mode == "precise" else 120
-
-    delta_quantised = np.round(delta_px / quantum) * quantum
-    for i in range(len(delta_quantised)):
-        if delta_quantised[i] == 0:
-            delta_quantised[i] = quantum * direction
-
-    # Adjust last event so total matches target distance
-    current_total = delta_quantised.sum()
-    diff = (distance * direction) - current_total
-    delta_quantised[-1] += diff
-
-    # Timestamps from log_dt
-    if len(log_dt) > 3:
-        median_log = float(np.median(log_dt))
-        log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
-        log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
-
-    dt_ms = np.exp(log_dt).clip(5, 80)
-
-    # Scale to realistic total duration
-    if mode == "fast":
-        expected_duration = distance * 0.2 + 100
-    elif mode == "precise":
-        expected_duration = distance * 1.5 + 300
-    else:
-        expected_duration = distance * 0.4 + 200
-
-    dt_ms = dt_ms * (expected_duration / max(dt_ms.sum(), 1.0))
-    dt_ms = dt_ms.clip(5, 80)
-
-    t_abs = np.cumsum(dt_ms).astype(int)
-    t_abs = np.concatenate([[0], t_abs[:-1]])
-
-    # Ensure monotonic
-    for i in range(1, len(t_abs)):
-        if t_abs[i] <= t_abs[i - 1]:
-            t_abs[i] = t_abs[i - 1] + 5
-
-    # Build events, removing zero-delta (keep at least 5)
-    events = []
-    for i in range(seq_len):
-        dy = int(delta_quantised[i])
-        if dy != 0 or len(events) < 5:
-            events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
-
-    return events

data/scroll_models/scroll_config.json

@@ -0,0 +1,13 @@
+{
+  "seq_len": 32,
+  "latent_dim": 16,
+  "hidden": 64,
+  "cond_dim": 7,
+  "epochs": 100,
+  "batch_size": 32,
+  "lr": 0.0005,
+  "beta_max": 0.3,
+  "beta_anneal_epochs": 15,
+  "w_delta": 1.0,
+  "w_logdt": 1.5
+}

examples/quickstart.py

@@ -0,0 +1,30 @@
+"""Minimal example: generate one trajectory + click event.
+
+Run:    uv run python examples/quickstart.py
+"""
+from __future__ import annotations
+
+from ai_mouse import generate
+
+start = (100, 200)
+end = (900, 400)
+points = generate(start, end, seed=0)
+
+print(f"Generated {len(points)} events:")
+print(f"  first move:  {points[0]}")
+print(f"  middle move: {points[len(points) // 2]}")
+print(f"  last move:   {points[-3]}")
+print(f"  click-down:  {points[-2]}")
+print(f"  click-up:    {points[-1]}")
+
+# Typical replay loop pattern. t_ms is cumulative from the start of the trace,
+# so block your sender thread until time-since-start reaches each event's t_ms.
+#
+# import time
+# t0 = time.monotonic()
+# for x, y, t_ms in points:
+#     target_wallclock = t0 + t_ms / 1000.0
+#     while time.monotonic() < target_wallclock:
+#         pass
+#     # replace this with pyautogui / pynput / win32 mouse_event:
+#     # send_mouse_move(x, y)

pyproject.toml

@@ -1,15 +1,37 @@
 [project]
 name = "ai-mouse"
-version = "0.1.0"
+version = "0.2.0"
+description = "Human-like mouse trajectory and scroll wheel event generator (ONNX Runtime SDK)."
 requires-python = ">=3.12,<3.14"
 dependencies = [
-    "torch>=2.2.0",
     "numpy>=1.26.0",
+    "onnxruntime>=1.17.0",
+]
+
+[project.urls]
+Repository = "https://github.com/<owner>/ai_mouse"
+
+[dependency-groups]
+dev = [
+    "torch>=2.2.0",
     "fastapi>=0.111.0",
     "uvicorn>=0.29.0",
     "scipy>=1.10.0",
     "matplotlib>=3.8.0",
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+    "httpx>=0.27.0",
+    "onnx>=1.15.0",
+    "onnxscript>=0.1",
 ]
 
-[dependency-groups]
+[build-system]
-dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0", "httpx>=0.27.0"]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ai_mouse"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]

src/ai_mouse/__init__.py

@@ -0,0 +1,78 @@
+"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events.
+
+Public API:
+
+    from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel
+
+See the project README for usage examples.
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Literal
+
+from ai_mouse import errors
+from ai_mouse._model_cache import get_mouse_model, get_scroll_model
+from ai_mouse.mouse import MouseModel
+from ai_mouse.scroll import ScrollModel
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "MouseModel",
+    "ScrollModel",
+    "errors",
+    "generate",
+    "generate_scroll",
+    "__version__",
+]
+
+
+def generate(
+    start: tuple[int, int],
+    end: tuple[int, int],
+    *,
+    n_points: int = 64,
+    speed: float | None = None,
+    click: bool = True,
+    seed: int | None = None,
+    model_path: str | Path | None = None,
+    providers: Sequence[str] | None = None,
+) -> list[tuple[int, int, int]]:
+    """Generate a human-like mouse trajectory.
+
+    See :meth:`MouseModel.generate` for argument semantics.
+    The underlying :class:`MouseModel` is cached process-wide; repeat
+    calls with the same ``(model_path, providers)`` reuse the session.
+    """
+    model = get_mouse_model(model_path, providers)
+    return model.generate(
+        start=start,
+        end=end,
+        n_points=n_points,
+        speed=speed,
+        click=click,
+        seed=seed,
+    )
+
+
+def generate_scroll(
+    start_scroll_y: int,
+    target_scroll_y: int,
+    *,
+    mode: Literal["target", "fast", "precise"] = "target",
+    viewport_height: int = 800,
+    seed: int | None = None,
+    model_path: str | Path | None = None,
+    providers: Sequence[str] | None = None,
+) -> list[dict]:
+    """Generate a sequence of mouse-wheel events. See :class:`ScrollModel.generate`."""
+    model = get_scroll_model(model_path, providers)
+    return model.generate(
+        start_scroll_y=start_scroll_y,
+        target_scroll_y=target_scroll_y,
+        mode=mode,
+        viewport_height=viewport_height,
+        seed=seed,
+    )

src/ai_mouse/_assets.py

@@ -0,0 +1,52 @@
+"""Asset path resolution for bundled ONNX weights and JSON metadata.
+
+Uses :mod:`importlib.resources` to locate files inside the installed
+package, falling back to a user-supplied directory if provided.
+"""
+from __future__ import annotations
+
+from importlib.resources import as_file, files
+from pathlib import Path
+
+from ai_mouse.errors import ModelLoadError
+
+_PACKAGE_ASSETS = "ai_mouse.assets"
+
+
+def bundled_path(name: str) -> Path:
+    """Return a filesystem path to a bundled asset.
+
+    Args:
+        name: filename inside the assets/ directory.
+
+    Returns:
+        A concrete :class:`pathlib.Path`. For non-zip wheels (the common
+        case) the path points directly into the installed package; for
+        zipapp installations :func:`importlib.resources.as_file`
+        materialises a temp file.
+    """
+    ref = files(_PACKAGE_ASSETS) / name
+    with as_file(ref) as p:
+        return Path(p)
+
+
+def resolve(model_path: Path | None, filename: str) -> Path:
+    """Locate an asset given an optional user-supplied directory.
+
+    Args:
+        model_path: user-supplied directory, or None to use bundled assets.
+        filename: file to locate inside the directory.
+
+    Returns:
+        Absolute path to the asset.
+
+    Raises:
+        ModelLoadError: if the file does not exist.
+    """
+    if model_path is None:
+        p = bundled_path(filename)
+    else:
+        p = Path(model_path) / filename
+    if not p.exists():
+        raise ModelLoadError(f"Required asset missing: {p}")
+    return p

src/ai_mouse/_model_cache.py

@@ -0,0 +1,45 @@
+"""Process-level lru_cache for default MouseModel / ScrollModel instances."""
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import lru_cache
+from pathlib import Path
+
+from ai_mouse.mouse import MouseModel
+from ai_mouse.scroll import ScrollModel
+
+
+@lru_cache(maxsize=4)
+def _get_mouse_model(
+    model_key: str,
+    providers_key: tuple[str, ...],
+) -> MouseModel:
+    path = None if model_key == "__bundled__" else Path(model_key)
+    providers = list(providers_key) if providers_key else None
+    return MouseModel(model_path=path, providers=providers)
+
+
+@lru_cache(maxsize=4)
+def _get_scroll_model(
+    model_key: str,
+    providers_key: tuple[str, ...],
+) -> ScrollModel:
+    path = None if model_key == "__bundled__" else Path(model_key)
+    providers = list(providers_key) if providers_key else None
+    return ScrollModel(model_path=path, providers=providers)
+
+
+def get_mouse_model(
+    model_path: str | Path | None,
+    providers: Sequence[str] | None,
+) -> MouseModel:
+    key = "__bundled__" if model_path is None else str(model_path)
+    return _get_mouse_model(key, tuple(providers or ()))
+
+
+def get_scroll_model(
+    model_path: str | Path | None,
+    providers: Sequence[str] | None,
+) -> ScrollModel:
+    key = "__bundled__" if model_path is None else str(model_path)
+    return _get_scroll_model(key, tuple(providers or ()))

src/ai_mouse/_postprocess.py

@@ -0,0 +1,180 @@
+"""Pure-numpy post-processing primitives for trajectory generation.
+
+All functions are pure (no I/O, no global state) and accept an explicit
+:class:`numpy.random.Generator` when randomness is involved.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
+    """5-tap gaussian smoothing along a 1-D array; endpoints preserved.
+
+    Args:
+        x: 1-D input array.
+        sigma: gaussian std. Default 1.0 gives weights approximately
+               [0.054, 0.244, 0.403, 0.244, 0.054].
+
+    Returns:
+        Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged.
+        If ``len(x) < 5`` returns a copy of ``x`` (kernel won't fit).
+    """
+    if len(x) < 5:
+        return x.copy()
+    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
+    kernel /= kernel.sum()
+    padded = np.pad(x, pad_width=2, mode="edge")
+    smoothed = np.convolve(padded, kernel, mode="valid")
+    smoothed[0] = x[0]
+    smoothed[-1] = x[-1]
+    return smoothed
+
+
+def snap_endpoints(
+    forward: np.ndarray,
+    lateral: np.ndarray,
+    seq_len: int,
+    n_snap: int = 6,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Force first point to (0,0) and last point to (1,0) with quadratic ease.
+
+    The last ``n_snap`` points are linearly interpolated towards (1, 0)
+    with quadratic easing, then the first/last points are pinned exactly.
+
+    Args:
+        forward: (T,) forward coordinates (modified in place).
+        lateral: (T,) lateral coordinates (modified in place).
+        seq_len: length of forward/lateral.
+        n_snap: number of trailing points to ease (capped at seq_len//4).
+
+    Returns:
+        ``(forward, lateral)`` after modification.
+    """
+    n_snap = min(n_snap, seq_len // 4)
+    for i in range(n_snap):
+        alpha = ((i + 1) / n_snap) ** 2
+        k = seq_len - n_snap + i
+        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
+        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
+    forward[0], lateral[0] = 0.0, 0.0
+    forward[-1], lateral[-1] = 1.0, 0.0
+    return forward, lateral
+
+
+def smooth_start(
+    forward: np.ndarray,
+    lateral: np.ndarray,
+    n: int = 4,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Dampen lateral oscillation in the first ``n`` points.
+
+    Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
+    forced non-decreasing locally; lateral is linearly damped towards 0.
+    """
+    n_start_fix = min(n, len(forward) // 4)
+    for i in range(1, n_start_fix + 1):
+        blend = i / (n_start_fix + 1)
+        forward[i] = max(forward[i], forward[i - 1])
+        lateral[i] = lateral[i] * blend
+    return forward, lateral
+
+
+def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
+    """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
+    seq_len = len(forward)
+    for i in range(1, seq_len - 1):
+        if forward[i] < forward[i - 1]:
+            forward[i] = forward[i - 1] + 0.001
+    forward = np.clip(forward, 0.0, 1.0)
+    forward[0] = 0.0
+    forward[-1] = 1.0
+    return forward
+
+
+def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
+    """Resample a 2-D polyline to ``n_points`` along cumulative arc length."""
+    arc = np.concatenate(
+        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
+    )
+    s_new = np.linspace(0, arc[-1], n_points)
+    return np.stack(
+        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
+        axis=1,
+    )
+
+
+def build_timestamps(
+    log_dt: np.ndarray,
+    total_duration_ms: float,
+    dt_clip: tuple[float, float] = (2.0, 150.0),
+) -> np.ndarray:
+    """Convert per-step log_dt + total duration to cumulative ms timestamps.
+
+    Args:
+        log_dt: (N,) array of natural-log step intervals.
+        total_duration_ms: target total span. The output is scaled so the
+            sum approximately matches this (modulo dt_clip).
+        dt_clip: (min, max) per-step clamp in milliseconds.
+
+    Returns:
+        (N,) integer-rounded cumulative timestamps starting at 0,
+        strictly increasing.
+    """
+    n = len(log_dt)
+    dt_raw = np.clip(np.exp(log_dt), 0.0, None)
+    dt_sum = dt_raw.sum()
+    if dt_sum > 1e-6:
+        scale = total_duration_ms / dt_sum
+    else:
+        scale = total_duration_ms / max(n, 1)
+    dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1])
+
+    t_abs = np.cumsum(dt_ms)
+    t_abs = np.concatenate([[0.0], t_abs[:-1]])
+
+    for i in range(1, n):
+        if t_abs[i] <= t_abs[i - 1]:
+            t_abs[i] = t_abs[i - 1] + 1.0
+    return t_abs
+
+
+def truncnorm_sample(
+    mu: float,
+    sigma: float,
+    low: float,
+    high: float,
+    rng: np.random.Generator,
+    max_tries: int = 32,
+) -> float:
+    """Sample from N(mu, sigma) truncated to [low, high] via rejection.
+
+    Falls back to clipping if rejection fails ``max_tries`` times.
+    """
+    for _ in range(max_tries):
+        v = rng.normal(mu, sigma)
+        if low <= v <= high:
+            return float(v)
+    return float(np.clip(rng.normal(mu, sigma), low, high))
+
+
+def sample_duration(
+    duration_dist: dict,
+    dist: float,
+    rng: np.random.Generator,
+) -> float:
+    """Sample total trajectory duration (ms) for the given pixel distance.
+
+    Uses per-bin log-normal parameters in ``duration_dist``.
+    """
+    bins = duration_dist["bins"]
+    params = duration_dist["params"]
+    bin_idx = len(bins) - 1
+    for i in range(len(bins) - 1):
+        if dist < bins[i + 1]:
+            bin_idx = i
+            break
+    bin_idx = min(bin_idx, len(params) - 1)
+    mu_log = params[bin_idx]["mu_log"]
+    sigma_log = params[bin_idx]["sigma_log"]
+    return float(np.exp(rng.normal(mu_log, sigma_log)))

src/ai_mouse/assets/click_dist.json

@@ -0,0 +1,6 @@
+{
+  "mu": 88.84602649006622,
+  "sigma": 22.60819023425299,
+  "low": 20.0,
+  "high": 500.0
+}

src/ai_mouse/assets/duration_dist.json

@@ -0,0 +1,47 @@
+{
+  "bins": [
+    0,
+    50,
+    100,
+    200,
+    400,
+    600,
+    800,
+    1200,
+    Infinity
+  ],
+  "params": [
+    {
+      "mu_log": 5.881647501269015,
+      "sigma_log": 0.05
+    },
+    {
+      "mu_log": 6.141460896343142,
+      "sigma_log": 0.43685929770118687
+    },
+    {
+      "mu_log": 6.409513572943365,
+      "sigma_log": 0.33676219820945824
+    },
+    {
+      "mu_log": 6.590028978462768,
+      "sigma_log": 0.2675760001657101
+    },
+    {
+      "mu_log": 6.70168704031046,
+      "sigma_log": 0.2809202882427188
+    },
+    {
+      "mu_log": 6.758762243710234,
+      "sigma_log": 0.21524718926818573
+    },
+    {
+      "mu_log": 6.214608098422191,
+      "sigma_log": 0.5
+    },
+    {
+      "mu_log": 6.214608098422191,
+      "sigma_log": 0.5
+    }
+  ]
+}

src/ai_mouse/assets/scroll_config.json

@@ -0,0 +1,13 @@
+{
+  "seq_len": 32,
+  "latent_dim": 16,
+  "hidden": 64,
+  "cond_dim": 7,
+  "epochs": 100,
+  "batch_size": 32,
+  "lr": 0.0005,
+  "beta_max": 0.3,
+  "beta_anneal_epochs": 15,
+  "w_delta": 1.0,
+  "w_logdt": 1.5
+}

src/ai_mouse/assets/train_config.json

@@ -0,0 +1,12 @@
+{
+  "seq_len": 64,
+  "epochs": 50,
+  "batch_size": 64,
+  "lr": 5e-06,
+  "d_model": 128,
+  "nhead": 4,
+  "num_layers": 4,
+  "dim_feedforward": 256,
+  "dropout": 0.1,
+  "cond_dim": 3
+}

src/ai_mouse/errors.py

@@ -0,0 +1,18 @@
+"""Exception hierarchy for the ai_mouse library.
+
+Downstream consumers can catch the umbrella :class:`AiMouseError`
+or the specific subclasses for finer control.
+"""
+from __future__ import annotations
+
+
+class AiMouseError(Exception):
+    """Base class for all ai_mouse errors."""
+
+
+class ModelLoadError(AiMouseError):
+    """Raised when ONNX weights / metadata cannot be loaded."""
+
+
+class GenerationError(AiMouseError):
+    """Raised when inference produces an invalid result (e.g. NaN)."""

src/ai_mouse/mouse.py

@@ -0,0 +1,166 @@
+"""MouseModel — ONNX Runtime-backed mouse trajectory generation."""
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import onnxruntime as ort
+
+from ai_mouse._assets import resolve
+from ai_mouse._coord import decode_trajectory
+from ai_mouse._postprocess import (
+    build_timestamps,
+    enforce_forward_monotonic,
+    gaussian_smooth,
+    resample_arc,
+    sample_duration,
+    smooth_start,
+    snap_endpoints,
+    truncnorm_sample,
+)
+from ai_mouse.errors import GenerationError, ModelLoadError
+
+_N_EULER_STEPS = 10
+
+
+class MouseModel:
+    """Persistent ONNX Runtime session for mouse trajectory generation."""
+
+    def __init__(
+        self,
+        model_path: str | Path | None = None,
+        providers: Sequence[str] | None = None,
+        seed: int | None = None,
+    ) -> None:
+        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None
+
+        onnx_path = resolve(path_obj, "flow_model.onnx")
+        cfg_path = resolve(path_obj, "train_config.json")
+        click_path = resolve(path_obj, "click_dist.json")
+        dur_path = resolve(path_obj, "duration_dist.json")
+
+        cfg = json.loads(cfg_path.read_text())
+        self._seq_len = int(cfg["seq_len"])
+        self._cond_dim = int(cfg.get("cond_dim", 3))
+        self._click_params = json.loads(click_path.read_text())
+        self._duration_dist = json.loads(dur_path.read_text())
+
+        try:
+            self._session = ort.InferenceSession(
+                str(onnx_path),
+                providers=list(providers) if providers else ["CPUExecutionProvider"],
+            )
+        except Exception as exc:
+            raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc
+
+        self._default_seed = seed
+        self._rng = np.random.default_rng(seed)
+
+    def generate(
+        self,
+        start: tuple[int, int],
+        end: tuple[int, int],
+        n_points: int = 64,
+        speed: float | None = None,
+        click: bool = True,
+        seed: int | None = None,
+    ) -> list[tuple[int, int, int]]:
+        rng = np.random.default_rng(seed) if seed is not None else self._rng
+
+        sx, sy = float(start[0]), float(start[1])
+        ex, ey = float(end[0]), float(end[1])
+        dist = max(math.hypot(ex - sx, ey - sy), 1.0)
+
+        total_duration = sample_duration(self._duration_dist, dist, rng)
+        if speed is not None and speed > 0:
+            total_duration /= speed
+        total_duration = max(total_duration, 10.0)
+
+        cond = np.array(
+            [
+                dist / 2000.0,
+                math.log(dist / 100.0),
+                math.log(total_duration / 500.0),
+            ],
+            dtype=np.float32,
+        )[None]
+
+        x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
+        dt = 1.0 / _N_EULER_STEPS
+        for step in range(_N_EULER_STEPS):
+            t = np.full((1,), step * dt, dtype=np.float32)
+            v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0]
+            x = x + v * dt
+
+        if not np.all(np.isfinite(x)):
+            raise GenerationError("Trajectory contains NaN/Inf after Euler integration")
+
+        forward = x[0, :, 0].copy()
+        lateral = x[0, :, 1].copy()
+        log_dt = x[0, :, 2].copy()
+
+        forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
+        forward, lateral = smooth_start(forward, lateral)
+        forward = enforce_forward_monotonic(forward)
+        lateral = gaussian_smooth(lateral, sigma=1.0)
+
+        log_dt = np.clip(log_dt, 0.0, 5.0)
+        log_dt[0] = 0.0
+
+        normalised = np.stack([forward, lateral], axis=1)
+        pixels = decode_trajectory(normalised, start, end)
+
+        if n_points != self._seq_len:
+            pixels = resample_arc(pixels, n_points)
+            log_dt = np.interp(
+                np.linspace(0, 1, n_points),
+                np.linspace(0, 1, self._seq_len),
+                log_dt,
+            )
+
+        ts = build_timestamps(log_dt, total_duration)
+
+        moves: list[tuple[int, int, int]] = [
+            (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i])))
+            for i in range(n_points)
+        ]
+        if not click:
+            return moves
+
+        click_dur = int(
+            truncnorm_sample(
+                float(self._click_params["mu"]),
+                float(self._click_params["sigma"]),
+                float(self._click_params["low"]),
+                float(self._click_params["high"]),
+                rng,
+            )
+        )
+        click_dur = max(click_dur, int(float(self._click_params["low"])))
+        last_t = moves[-1][2]
+        cx, cy = moves[-1][0], moves[-1][1]
+        return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)]
+
+    def sample_click_duration_ms(self, seed: int | None = None) -> int:
+        rng = np.random.default_rng(seed) if seed is not None else self._rng
+        v = truncnorm_sample(
+            float(self._click_params["mu"]),
+            float(self._click_params["sigma"]),
+            float(self._click_params["low"]),
+            float(self._click_params["high"]),
+            rng,
+        )
+        return max(int(v), int(float(self._click_params["low"])))
+
+    def close(self) -> None:
+        self._session = None  # type: ignore[assignment]
+
+    def __enter__(self) -> "MouseModel":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

src/ai_mouse/scroll.py

@@ -0,0 +1,139 @@
+"""ScrollModel — ONNX Runtime-backed scroll event generation."""
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Literal, Optional
+
+import numpy as np
+import onnxruntime as ort
+
+from ai_mouse._assets import resolve
+from ai_mouse.errors import ModelLoadError
+
+_DURATION_TABLE = {
+    "fast":    lambda d: d * 0.2 + 100.0,
+    "precise": lambda d: d * 1.5 + 300.0,
+    "target":  lambda d: d * 0.4 + 200.0,
+}
+
+_QUANTUM = {"precise": 40, "fast": 120, "target": 120}
+
+
+class ScrollModel:
+    """Persistent ONNX Runtime session for scroll event generation."""
+
+    def __init__(
+        self,
+        model_path: str | Path | None = None,
+        providers: Sequence[str] | None = None,
+        seed: int | None = None,
+    ) -> None:
+        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None
+
+        onnx_path = resolve(path_obj, "scroll_decoder.onnx")
+        cfg_path = resolve(path_obj, "scroll_config.json")
+        cfg = json.loads(cfg_path.read_text())
+
+        self._seq_len = int(cfg["seq_len"])
+        self._latent_dim = int(cfg["latent_dim"])
+        self._cond_dim = int(cfg["cond_dim"])
+
+        try:
+            self._session = ort.InferenceSession(
+                str(onnx_path),
+                providers=list(providers) if providers else ["CPUExecutionProvider"],
+            )
+        except Exception as exc:
+            raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc
+
+        self._rng = np.random.default_rng(seed)
+
+    def generate(
+        self,
+        start_scroll_y: int,
+        target_scroll_y: int,
+        mode: Literal["target", "fast", "precise"] = "target",
+        viewport_height: int = 800,
+        seed: int | None = None,
+    ) -> list[dict]:
+        rng = np.random.default_rng(seed) if seed is not None else self._rng
+
+        distance = abs(target_scroll_y - start_scroll_y)
+        direction = 1 if target_scroll_y > start_scroll_y else -1
+        distance = max(distance, 10)
+
+        cond = self._build_condition(float(distance), direction, mode, viewport_height)
+        z = rng.standard_normal((1, self._latent_dim)).astype(np.float32)
+        decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0]
+
+        delta_norm = decoded[:, 0]
+        log_dt = decoded[:, 1]
+
+        delta_weights = np.exp(delta_norm)
+        delta_weights /= delta_weights.sum()
+        delta_px = delta_weights * distance * direction
+
+        quantum = _QUANTUM[mode]
+        delta_q = np.round(delta_px / quantum) * quantum
+        for i in range(len(delta_q)):
+            if delta_q[i] == 0:
+                delta_q[i] = quantum * direction
+        delta_q[-1] += (distance * direction) - delta_q.sum()
+
+        if len(log_dt) > 3:
+            median_log = float(np.median(log_dt))
+            log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
+            log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
+        dt_ms = np.clip(np.exp(log_dt), 5, 80)
+        expected = _DURATION_TABLE[mode](distance)
+        dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80)
+
+        t_abs = np.cumsum(dt_ms).astype(int)
+        t_abs = np.concatenate([[0], t_abs[:-1]])
+        for i in range(1, len(t_abs)):
+            if t_abs[i] <= t_abs[i - 1]:
+                t_abs[i] = t_abs[i - 1] + 5
+
+        events: list[dict] = []
+        for i in range(self._seq_len):
+            dy = int(delta_q[i])
+            if dy != 0 or len(events) < 5:
+                events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
+        return events
+
+    def _build_condition(
+        self,
+        distance: float,
+        direction: int,
+        mode: str,
+        viewport_height: int,
+    ) -> np.ndarray:
+        mode_onehot = [0.0, 0.0, 0.0]
+        if mode == "target":
+            mode_onehot[0] = 1.0
+        elif mode == "fast":
+            mode_onehot[1] = 1.0
+        elif mode == "precise":
+            mode_onehot[2] = 1.0
+        return np.array(
+            [
+                distance / 5000.0,
+                math.log(max(distance, 1.0) / 500.0),
+                float(direction),
+                viewport_height / 1000.0,
+                *mode_onehot,
+            ],
+            dtype=np.float32,
+        )
+
+    def close(self) -> None:
+        self._session = None  # type: ignore[assignment]
+
+    def __enter__(self) -> "ScrollModel":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

tests/test_coord.py

@@ -1,113 +0,0 @@
-"""Tests for rotated coordinate system transforms."""
-from __future__ import annotations
-
-import math
-
-import numpy as np
-import pytest
-
-from ai_mouse.coord import encode_trajectory, decode_trajectory
-
-
-class TestEncodeTrajectory:
-    """Test pixel → rotated normalised frame."""
-
-    def test_start_maps_to_origin(self):
-        start = (100, 200)
-        end = (400, 500)
-        points = np.array([[100, 200]], dtype=float)
-        result = encode_trajectory(points, start, end)
-        np.testing.assert_allclose(result[0], [0.0, 0.0], atol=1e-10)
-
-    def test_end_maps_to_one_zero(self):
-        start = (100, 200)
-        end = (400, 500)
-        points = np.array([[400, 500]], dtype=float)
-        result = encode_trajectory(points, start, end)
-        np.testing.assert_allclose(result[0], [1.0, 0.0], atol=1e-10)
-
-    def test_midpoint_maps_to_half_zero(self):
-        start = (0, 0)
-        end = (200, 0)
-        points = np.array([[100, 0]], dtype=float)
-        result = encode_trajectory(points, start, end)
-        np.testing.assert_allclose(result[0], [0.5, 0.0], atol=1e-10)
-
-    def test_lateral_offset_positive(self):
-        """Point at (100, 50) with horizontal start→end has lateral = 50/200 = 0.25."""
-        start = (0, 0)
-        end = (200, 0)
-        # For horizontal u=(1,0), v=(-0, 1)=(0,1).
-        # Point (100, 50): forward = 100/200=0.5, lateral = 50/200=0.25
-        points = np.array([[100, 50]], dtype=float)
-        result = encode_trajectory(points, start, end)
-        np.testing.assert_allclose(result[0], [0.5, 0.25], atol=1e-10)
-
-    def test_various_angles(self):
-        """Encode/decode round-trip works for various angles."""
-        angles = [0, 45, 90, 135, 180, -45, -90, -135]
-        for deg in angles:
-            rad = math.radians(deg)
-            start = (400, 300)
-            dist = 200
-            end = (int(400 + dist * math.cos(rad)), int(300 + dist * math.sin(rad)))
-            # Create a curved path
-            t = np.linspace(0, 1, 20)
-            px = start[0] + t * (end[0] - start[0]) + 20 * np.sin(t * math.pi)
-            py = start[1] + t * (end[1] - start[1]) + 20 * np.cos(t * math.pi)
-            points = np.stack([px, py], axis=1)
-
-            encoded = encode_trajectory(points, start, end)
-            assert encoded[0, 0] == pytest.approx(0.0, abs=0.2)
-            assert encoded[-1, 0] == pytest.approx(1.0, abs=0.2)
-
-
-class TestDecodeTrajectory:
-    """Test rotated normalised frame → pixel."""
-
-    def test_origin_maps_to_start(self):
-        start = (100, 200)
-        end = (400, 500)
-        normalised = np.array([[0.0, 0.0]], dtype=float)
-        result = decode_trajectory(normalised, start, end)
-        np.testing.assert_allclose(result[0], [100, 200], atol=1e-10)
-
-    def test_one_zero_maps_to_end(self):
-        start = (100, 200)
-        end = (400, 500)
-        normalised = np.array([[1.0, 0.0]], dtype=float)
-        result = decode_trajectory(normalised, start, end)
-        np.testing.assert_allclose(result[0], [400, 500], atol=1e-10)
-
-
-class TestRoundTrip:
-    """Encode then decode should return original points."""
-
-    def test_round_trip_horizontal(self):
-        start = (50, 100)
-        end = (350, 100)
-        points = np.array([[50, 100], [150, 130], [250, 90], [350, 100]], dtype=float)
-        encoded = encode_trajectory(points, start, end)
-        decoded = decode_trajectory(encoded, start, end)
-        np.testing.assert_allclose(decoded, points, atol=1e-8)
-
-    def test_round_trip_diagonal(self):
-        start = (100, 100)
-        end = (500, 400)
-        rng = np.random.default_rng(42)
-        points = np.column_stack([
-            np.linspace(100, 500, 30) + rng.normal(0, 10, 30),
-            np.linspace(100, 400, 30) + rng.normal(0, 10, 30),
-        ])
-        encoded = encode_trajectory(points, start, end)
-        decoded = decode_trajectory(encoded, start, end)
-        np.testing.assert_allclose(decoded, points, atol=1e-8)
-
-    def test_round_trip_vertical(self):
-        """Vertical movement (angle=90°) doesn't collapse."""
-        start = (300, 50)
-        end = (300, 450)
-        points = np.array([[300, 50], [310, 200], [295, 350], [300, 450]], dtype=float)
-        encoded = encode_trajectory(points, start, end)
-        decoded = decode_trajectory(encoded, start, end)
-        np.testing.assert_allclose(decoded, points, atol=1e-8)

tests/test_generator.py

@@ -1,158 +0,0 @@
-"""Tests for Flow Matching trajectory generator."""
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from ai_mouse.generator import generate
-from ai_mouse.models import TrajectoryFlowModel
-
-
-@pytest.fixture
-def model_dir(tmp_path):
-    """Create temp dir with Flow model artifacts."""
-    model = TrajectoryFlowModel(seq_len=64)
-    torch.save(model.state_dict(), tmp_path / "flow_model.pt")
-
-    click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
-    (tmp_path / "click_dist.json").write_text(json.dumps(click_dist))
-
-    duration_dist = {
-        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
-        "params": [
-            {"mu_log": 5.5, "sigma_log": 0.3},
-            {"mu_log": 5.8, "sigma_log": 0.3},
-            {"mu_log": 6.0, "sigma_log": 0.3},
-            {"mu_log": 6.2, "sigma_log": 0.3},
-            {"mu_log": 6.5, "sigma_log": 0.3},
-            {"mu_log": 6.7, "sigma_log": 0.3},
-            {"mu_log": 6.9, "sigma_log": 0.3},
-            {"mu_log": 7.0, "sigma_log": 0.3},
-        ],
-    }
-    (tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist))
-
-    train_config = {
-        "seq_len": 64,
-        "d_model": 128,
-        "nhead": 4,
-        "num_layers": 4,
-        "dim_feedforward": 256,
-        "cond_dim": 3,
-    }
-    (tmp_path / "train_config.json").write_text(json.dumps(train_config))
-
-    return tmp_path
-
-
-class TestGenerate:
-    def test_returns_list_of_tuples(self, model_dir):
-        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-        assert isinstance(result, list)
-        assert all(isinstance(p, tuple) and len(p) == 3 for p in result)
-        # All elements are ints
-        for p in result:
-            assert all(isinstance(v, int) for v in p)
-
-    def test_timestamps_monotonically_increasing(self, model_dir):
-        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-        times = [p[2] for p in result]
-        for i in range(1, len(times)):
-            assert times[i] >= times[i - 1]
-
-    def test_starts_near_start(self, model_dir):
-        start = (100, 200)
-        result = generate(start=start, end=(500, 400), model_dir=str(model_dir))
-        first = result[0]
-        assert abs(first[0] - start[0]) < 30
-        assert abs(first[1] - start[1]) < 30
-
-    def test_ends_near_end(self, model_dir):
-        end = (500, 400)
-        result = generate(start=(100, 200), end=end, model_dir=str(model_dir))
-        # Last two are click events; the one before is last movement point
-        last_move = result[-3]
-        assert abs(last_move[0] - end[0]) < 30
-        assert abs(last_move[1] - end[1]) < 30
-
-    def test_last_two_are_click_events(self, model_dir):
-        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-        down = result[-2]
-        up = result[-1]
-        # Same x, y for click down and up
-        assert down[0] == up[0]
-        assert down[1] == up[1]
-        # Up timestamp > down timestamp
-        assert up[2] > down[2]
-        # Click duration within bounds
-        assert 20 <= up[2] - down[2] <= 300
-
-    def test_different_z_gives_different_paths(self, model_dir):
-        r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-        r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-        points1 = [(p[0], p[1]) for p in r1[:-2]]
-        points2 = [(p[0], p[1]) for p in r2[:-2]]
-        assert points1 != points2
-
-    def test_n_points_parameter(self, model_dir):
-        result = generate(
-            start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir)
-        )
-        # 32 move points + 2 click events = 34
-        assert len(result) == 34
-
-
-class TestPostProcessing:
-    def test_dt_diversity_preserved(self, model_dir):
-        """After removing speed_profile + median clip, multiple generations
-        should differ in their Δt sequences (not all identical)."""
-        results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
-                   for _ in range(5)]
-        # Extract Δt sequences (only move events, not click events)
-        dts = []
-        for r in results:
-            moves = r[:-2]
-            dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)]
-            dts.append(dt_seq)
-        # At least 2 of the 5 sequences should differ at any given index
-        for i in range(min(len(d) for d in dts)):
-            values = {tuple([d[i]]) for d in dts}
-            if len(values) > 1:
-                return  # at least one position has variation — pass
-        pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")
-
-
-class TestGaussianSmooth:
-    def test_endpoints_preserved(self):
-        from ai_mouse.generator import _gaussian_smooth
-        x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
-        smoothed = _gaussian_smooth(x, sigma=1.0)
-        assert smoothed[0] == 1.0
-        assert smoothed[-1] == 2.0
-
-    def test_smooths_high_frequency(self):
-        """A high-frequency square wave should have reduced amplitude after smoothing."""
-        from ai_mouse.generator import _gaussian_smooth
-        x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
-        smoothed = _gaussian_smooth(x, sigma=1.0)
-        # Interior amplitude should be reduced
-        interior_orig = x[2:-2]
-        interior_smooth = smoothed[2:-2]
-        assert interior_smooth.std() < interior_orig.std()
-
-    def test_constant_signal_unchanged(self):
-        from ai_mouse.generator import _gaussian_smooth
-        x = np.full(20, 0.5, dtype=np.float64)
-        smoothed = _gaussian_smooth(x, sigma=1.0)
-        np.testing.assert_allclose(smoothed, x, rtol=1e-6)
-
-    def test_short_array_returns_unchanged(self):
-        """Arrays shorter than the kernel are returned unchanged."""
-        from ai_mouse.generator import _gaussian_smooth
-        x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
-        smoothed = _gaussian_smooth(x, sigma=1.0)
-        np.testing.assert_allclose(smoothed, x, rtol=1e-6)

tests/test_scroll_generator.py

@@ -1,50 +0,0 @@
-"""Tests for scroll generator."""
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import torch
-import pytest
-
-from ai_mouse.scroll.generator import generate_scroll
-from ai_mouse.scroll.models import ScrollCVAE
-
-
-@pytest.fixture
-def scroll_model_dir(tmp_path):
-    model = ScrollCVAE(seq_len=32)
-    torch.save(model.state_dict(), tmp_path / "scroll_model.pt")
-    config = {"seq_len": 32, "epochs": 100}
-    (tmp_path / "scroll_config.json").write_text(json.dumps(config))
-    return tmp_path
-
-
-class TestGenerateScroll:
-    def test_returns_list_of_dicts(self, scroll_model_dir):
-        result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
-        assert isinstance(result, list)
-        assert len(result) > 0
-        assert all("deltaY" in e and "t" in e and "deltaMode" in e for e in result)
-
-    def test_timestamps_monotonic(self, scroll_model_dir):
-        result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
-        times = [e["t"] for e in result]
-        for i in range(1, len(times)):
-            assert times[i] >= times[i - 1]
-
-    def test_total_scroll_approximately_matches_distance(self, scroll_model_dir):
-        result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
-        total = sum(e["deltaY"] for e in result)
-        # Should be within 30% of target distance (2000px)
-        assert abs(total - 2000) < 2000 * 0.4
-
-    def test_deltaY_are_integers(self, scroll_model_dir):
-        result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
-        assert all(isinstance(e["deltaY"], int) for e in result)
-
-    def test_direction_up(self, scroll_model_dir):
-        result = generate_scroll(3000, 1000, mode="target", model_dir=str(scroll_model_dir))
-        total = sum(e["deltaY"] for e in result)
-        # Negative total for scrolling up
-        assert total < 0

tests/tools/conftest.py

@@ -8,8 +8,8 @@ import numpy as np
 import pytest
 import torch
 
-from ai_mouse.models import TrajectoryFlowModel
+from tools.models import TrajectoryFlowModel
-from ai_mouse.scroll.models import ScrollCVAE
+from tools.scroll.models import ScrollCVAE
 
 
 @pytest.fixture

tests/tools/test_balabit_adapter.py

@@ -9,7 +9,7 @@ import pytest
 
 def test_module_exports():
     """The adapter module must export the public functions used by CLI."""
-    from ai_mouse.data_adapters import balabit
+    from tools.data_adapters import balabit
     assert hasattr(balabit, "parse_session_csv")
     assert hasattr(balabit, "segment_by_clicks")
     assert hasattr(balabit, "filter_segments")
@@ -20,7 +20,7 @@ def test_module_exports():
 
 def test_mouse_event_dataclass():
     """MouseEvent has expected fields."""
-    from ai_mouse.data_adapters.balabit import MouseEvent
+    from tools.data_adapters.balabit import MouseEvent
     e = MouseEvent(t_ms=100, button="NoButton", state="Move", x=300, y=400)
     assert e.t_ms == 100
     assert e.state == "Move"
@@ -29,7 +29,7 @@ def test_mouse_event_dataclass():
 
 def test_segment_dataclass():
     """Segment has expected fields."""
-    from ai_mouse.data_adapters.balabit import MouseEvent, Segment
+    from tools.data_adapters.balabit import MouseEvent, Segment
     events = [MouseEvent(t_ms=0, button="NoButton", state="Move", x=10, y=20)]
     s = Segment(events=events, click_x=100, click_y=200, click_t_ms=500, session_id="user1_s1")
     assert s.events == events
@@ -45,7 +45,7 @@ def _write_csv(path: Path, rows: list[str]) -> None:
 
 class TestParseSessionCsv:
     def test_parses_basic_rows(self, tmp_path):
-        from ai_mouse.data_adapters.balabit import parse_session_csv
+        from tools.data_adapters.balabit import parse_session_csv
         csv = tmp_path / "session_1"
         _write_csv(csv, [
             "1500000000.000,0.000,NoButton,Move,100,200",
@@ -63,7 +63,7 @@ class TestParseSessionCsv:
 
     def test_handles_float_timestamps(self, tmp_path):
         """Client timestamps are floats in seconds; we convert to int ms."""
-        from ai_mouse.data_adapters.balabit import parse_session_csv
+        from tools.data_adapters.balabit import parse_session_csv
         csv = tmp_path / "session_2"
         _write_csv(csv, [
             "0,1.234,NoButton,Move,50,60",
@@ -75,7 +75,7 @@ class TestParseSessionCsv:
 
     def test_skips_malformed_rows(self, tmp_path):
         """Rows with bad data are logged and skipped, not raised."""
-        from ai_mouse.data_adapters.balabit import parse_session_csv
+        from tools.data_adapters.balabit import parse_session_csv
         csv = tmp_path / "session_3"
         _write_csv(csv, [
             "0,0.000,NoButton,Move,100,200",
@@ -89,7 +89,7 @@ class TestParseSessionCsv:
         assert events[1].x == 150
 
     def test_returns_empty_list_for_empty_file(self, tmp_path):
-        from ai_mouse.data_adapters.balabit import parse_session_csv
+        from tools.data_adapters.balabit import parse_session_csv
         csv = tmp_path / "session_4"
         csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8")
         events = parse_session_csv(csv)
@@ -98,11 +98,11 @@ class TestParseSessionCsv:
 
 class TestSegmentByClicks:
     def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"):
-        from ai_mouse.data_adapters.balabit import MouseEvent
+        from tools.data_adapters.balabit import MouseEvent
         return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y)
 
     def test_one_click_one_segment(self):
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),
             self._make_event(100, "Move", 50, 60),
@@ -120,7 +120,7 @@ class TestSegmentByClicks:
 
     def test_window_excludes_old_events(self):
         """Move events earlier than (click_t - window_ms) are dropped."""
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),       # too old
             self._make_event(100, "Move", 20, 30),     # too old
@@ -133,7 +133,7 @@ class TestSegmentByClicks:
         assert segments[0].events[0].t_ms == 900
 
     def test_multiple_clicks_multiple_segments(self):
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),
             self._make_event(100, "Pressed", 50, 50, button="Left"),
@@ -152,7 +152,7 @@ class TestSegmentByClicks:
 
     def test_skips_pressed_with_non_left_button(self):
         """Right-clicks and wheel-clicks don't anchor segments (only Left)."""
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),
             self._make_event(100, "Pressed", 50, 50, button="Right"),  # ignored
@@ -164,7 +164,7 @@ class TestSegmentByClicks:
         assert segments[0].click_x == 70
 
     def test_no_clicks_returns_empty(self):
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),
             self._make_event(100, "Move", 20, 30),
@@ -174,7 +174,7 @@ class TestSegmentByClicks:
 
     def test_excludes_drag_events(self):
         """Drag events are not Move; segment should only include Move."""
-        from ai_mouse.data_adapters.balabit import segment_by_clicks
+        from tools.data_adapters.balabit import segment_by_clicks
         events = [
             self._make_event(0, "Move", 10, 20),
             self._make_event(100, "Drag", 30, 40),       # not Move
@@ -191,7 +191,7 @@ class TestSegmentByClicks:
 class TestFilterSegments:
     def _seg(self, events_data: list[tuple[int, int, int]], click=(500, 500), session="s"):
         """events_data: list of (t_ms, x, y) tuples."""
-        from ai_mouse.data_adapters.balabit import MouseEvent, Segment
+        from tools.data_adapters.balabit import MouseEvent, Segment
         events = [MouseEvent(t_ms=t, button="NoButton", state="Move", x=x, y=y)
                   for (t, x, y) in events_data]
         click_t = events[-1].t_ms + 50 if events else 100
@@ -199,14 +199,14 @@ class TestFilterSegments:
                        click_t_ms=click_t, session_id=session)
 
     def test_drops_segment_with_too_few_events(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         seg = self._seg([(0, 100, 100), (50, 105, 105)])  # 2 events, min=5
         result = filter_segments([seg], min_events=5, min_dist=10,
                                   max_span_ms=5000, max_gap_ms=200)
         assert result == []
 
     def test_drops_segment_with_short_distance(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # Start (100,100), end click=(105,100) → dist=5, min_dist=50
         events = [(i*10, 100+i, 100) for i in range(10)]
         seg = self._seg(events, click=(105, 100))
@@ -215,7 +215,7 @@ class TestFilterSegments:
         assert result == []
 
     def test_drops_segment_with_too_long_span(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # Span 6000ms, max_span=5000
         events = [(0, 100, 100), (2000, 200, 200), (4000, 300, 300), (6000, 400, 400),
                   (6010, 410, 400), (6020, 420, 400)]
@@ -225,7 +225,7 @@ class TestFilterSegments:
         assert result == []
 
     def test_drops_segment_with_gap(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # Gap of 500ms between events 2 and 3, max_gap=200
         events = [(0, 100, 100), (50, 110, 110), (100, 120, 120),
                   (600, 200, 200), (650, 210, 210), (700, 220, 220)]
@@ -235,7 +235,7 @@ class TestFilterSegments:
         assert result == []
 
     def test_drops_segment_with_out_of_range_coords(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # x=10000 out of range
         events = [(i*10, 10000, 100) for i in range(6)]
         seg = self._seg(events, click=(10100, 100))
@@ -245,7 +245,7 @@ class TestFilterSegments:
 
     def test_drops_segment_with_short_arc_length(self):
         """Total arc < 50 even though endpoints are far → high-frequency jitter only."""
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # Tiny back-and-forth with click 100px away — unrealistic, drop it
         events = [(i*10, 100 + (i % 2), 100) for i in range(10)]
         seg = self._seg(events, click=(200, 100))
@@ -255,7 +255,7 @@ class TestFilterSegments:
         assert result == []
 
     def test_keeps_valid_segment(self):
-        from ai_mouse.data_adapters.balabit import filter_segments
+        from tools.data_adapters.balabit import filter_segments
         # Smooth 100→500 px straight line, 10 events, span 500ms
         events = [(i*50, 100 + i*40, 100) for i in range(10)]
         seg = self._seg(events, click=(500, 100))
@@ -266,8 +266,8 @@ class TestFilterSegments:
 
 class TestProcessSession:
     def test_writes_jsonl_in_expected_format(self, tmp_path):
-        from ai_mouse.config import BalabitAdapterConfig
+        from tools.config import BalabitAdapterConfig
-        from ai_mouse.data_adapters.balabit import process_session
+        from tools.data_adapters.balabit import process_session
 
         # Construct a Balabit-format CSV with one valid segment
         csv_path = tmp_path / "user_session_42"
@@ -303,8 +303,8 @@ class TestProcessSession:
         assert record["events"][0]["t"] == 0
 
     def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
-        from ai_mouse.config import BalabitAdapterConfig
+        from tools.config import BalabitAdapterConfig
-        from ai_mouse.data_adapters.balabit import process_session
+        from tools.data_adapters.balabit import process_session
 
         csv_path = tmp_path / "empty_session"
         _write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"])  # no clicks
@@ -318,8 +318,8 @@ class TestProcessSession:
             assert out.read_text() == ""
 
     def test_appends_to_existing_jsonl(self, tmp_path):
-        from ai_mouse.config import BalabitAdapterConfig
+        from tools.config import BalabitAdapterConfig
-        from ai_mouse.data_adapters.balabit import process_session
+        from tools.data_adapters.balabit import process_session
 
         out = tmp_path / "out.jsonl"
         out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")

tests/tools/test_eval_metrics.py

@@ -8,7 +8,7 @@ import pytest
 class TestKinematics:
     def test_compute_speed_constant_velocity(self):
         """Constant-velocity trajectory has constant speed."""
-        from ai_mouse.eval.metrics import compute_speed
+        from tools.eval.metrics import compute_speed
         # 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms
         xs = np.arange(0, 100, 10, dtype=float)
         ys = np.zeros(10, dtype=float)
@@ -20,7 +20,7 @@ class TestKinematics:
 
     def test_compute_speed_handles_zero_dt(self):
         """Adjacent points with same timestamp must not produce NaN/inf."""
-        from ai_mouse.eval.metrics import compute_speed
+        from tools.eval.metrics import compute_speed
         xs = np.array([0.0, 10.0, 20.0])
         ys = np.array([0.0, 0.0, 0.0])
         ts = np.array([0.0, 0.0, 100.0])  # zero dt between [0] and [1]
@@ -29,7 +29,7 @@ class TestKinematics:
 
     def test_compute_acceleration(self):
         """Linearly increasing speed → constant acceleration."""
-        from ai_mouse.eval.metrics import compute_acceleration
+        from tools.eval.metrics import compute_acceleration
         # speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms²
         speeds = np.array([0.1, 0.2, 0.3, 0.4])
         ts = np.array([100.0, 200.0, 300.0, 400.0])
@@ -37,7 +37,7 @@ class TestKinematics:
         np.testing.assert_allclose(a, 0.001, rtol=1e-4)
 
     def test_compute_jerk(self):
-        from ai_mouse.eval.metrics import compute_jerk
+        from tools.eval.metrics import compute_jerk
         # accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001
         accels = np.array([0.001, 0.002, 0.003])
         ts = np.array([200.0, 300.0, 400.0])
@@ -47,7 +47,7 @@ class TestKinematics:
 
 class TestStatsSummary:
     def test_compute_stats_returns_expected_keys(self):
-        from ai_mouse.eval.metrics import compute_stats
+        from tools.eval.metrics import compute_stats
         x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
         s = compute_stats(x)
         assert "mean" in s
@@ -59,7 +59,7 @@ class TestStatsSummary:
         assert "p95" in s
 
     def test_cv_for_constant_is_zero(self):
-        from ai_mouse.eval.metrics import compute_stats
+        from tools.eval.metrics import compute_stats
         x = np.full(10, 3.0)
         s = compute_stats(x)
         assert s["cv"] == 0.0
@@ -68,7 +68,7 @@ class TestStatsSummary:
 class TestFftSpectrum:
     def test_finds_dominant_frequency(self):
         """A pure 8 Hz signal should have its peak near 8 Hz."""
-        from ai_mouse.eval.metrics import fft_spectrum
+        from tools.eval.metrics import fft_spectrum
         # Sample at 100 Hz for 1 second
         sample_rate_hz = 100.0
         ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz)
@@ -78,7 +78,7 @@ class TestFftSpectrum:
         assert abs(peak_freq - 8.0) < 1.0  # within 1 Hz
 
     def test_returns_only_positive_frequencies(self):
-        from ai_mouse.eval.metrics import fft_spectrum
+        from tools.eval.metrics import fft_spectrum
         signal = np.random.randn(64)
         freqs, mags = fft_spectrum(signal, 50.0)
         assert (freqs >= 0).all()
@@ -88,7 +88,7 @@ class TestFftSpectrum:
 class TestKlDivergence:
     def test_identical_distributions_zero_kl(self):
         """KL(p, p) ≈ 0."""
-        from ai_mouse.eval.metrics import kl_divergence_histograms
+        from tools.eval.metrics import kl_divergence_histograms
         rng = np.random.default_rng(42)
         x = rng.normal(0, 1, 5000)
         y = rng.normal(0, 1, 5000)
@@ -97,7 +97,7 @@ class TestKlDivergence:
 
     def test_different_distributions_positive_kl(self):
         """Different means → positive KL."""
-        from ai_mouse.eval.metrics import kl_divergence_histograms
+        from tools.eval.metrics import kl_divergence_histograms
         rng = np.random.default_rng(42)
         x = rng.normal(0, 1, 5000)
         y = rng.normal(3, 1, 5000)
@@ -106,7 +106,7 @@ class TestKlDivergence:
 
     def test_handles_disjoint_supports(self):
         """No NaN even when histograms have non-overlapping bins."""
-        from ai_mouse.eval.metrics import kl_divergence_histograms
+        from tools.eval.metrics import kl_divergence_histograms
         x = np.array([1.0, 1.1, 1.2, 1.3, 1.4])
         y = np.array([10.0, 10.1, 10.2, 10.3, 10.4])
         kl = kl_divergence_histograms(x, y, bins=10)
@@ -116,7 +116,7 @@ class TestKlDivergence:
 class TestReportGeneration:
     def test_generates_report_md(self, tmp_path):
         """Smoke test: build_report writes an MD file with all expected sections."""
-        from ai_mouse.eval.report import build_report
+        from tools.eval.report import build_report
 
         # Synthetic generated traces (3 traces, 50 points each)
         rng = np.random.default_rng(0)

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)

tests/tools/test_models.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 import torch
 import pytest
 
-from ai_mouse.models import TrajectoryFlowModel
+from tools.models import TrajectoryFlowModel
 
 
 class TestTrajectoryFlowModel:

tests/tools/test_scroll_collector.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import pytest
 
-from ai_mouse.scroll.collector import ScrollCollector
+from tools.scroll.collector import ScrollCollector
 
 
 class TestNextTarget:

tests/tools/test_scroll_models.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 import torch
 import pytest
 
-from ai_mouse.scroll.models import ScrollCVAE
+from tools.scroll.models import ScrollCVAE
 
 
 class TestScrollCVAEForward:

tests/tools/test_scroll_trainer.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import numpy as np
 import pytest
 
-from ai_mouse.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll
+from tools.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll
 
 
 def _make_synthetic_scroll_trace(mode="target"):

tests/tools/test_server.py

@@ -7,8 +7,8 @@ import pytest
 import pytest_asyncio
 from httpx import ASGITransport, AsyncClient
 
-from ai_mouse.server import create_app
+from tools.server import create_app
-from ai_mouse.server.deps import get_data_dir
+from tools.server.deps import get_data_dir
 
 
 @pytest.fixture
@@ -87,7 +87,7 @@ class TestCollect:
     async def test_collect_trace_increments_count(self, client, tmp_path, monkeypatch):
         """Test that posting a trace increments the collected count."""
         # Monkeypatch data dir to use tmp
-        import ai_mouse.server.deps as deps
+        import tools.server.deps as deps
         monkeypatch.setattr(deps, "_DATA_DIR", tmp_path)
 
         # Start collection
@@ -124,7 +124,7 @@ class TestVerify:
     @pytest.mark.asyncio
     async def test_verify_returns_paths(self, client, model_dir, monkeypatch):
         """Test trajectory generation endpoint."""
-        import ai_mouse.server.routes_verify as rv
+        import tools.server.routes_verify as rv
         # We can't easily monkeypatch the model dir used inside the route
         # but we can test the endpoint is accessible
         resp = await client.post(

tests/tools/test_trainer.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import numpy as np
 import pytest
 
-from ai_mouse.trainer import load_and_prepare_data, train, _augment
+from tools.trainer import load_and_prepare_data, train, _augment
 
 
 def _make_synthetic_trace(start, end, n_moves=30):
@@ -124,21 +124,21 @@ class TestTrain:
 class TestTrajectoryDataset:
     def test_dataset_length_with_augmentation(self):
         """Dataset length = N * 6 when augment=True."""
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((10, 64, 3), dtype=np.float32)
         cond = np.zeros((10, 3), dtype=np.float32)
         ds = TrajectoryDataset(seq, cond, augment=True)
         assert len(ds) == 60
 
     def test_dataset_length_without_augmentation(self):
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((10, 64, 3), dtype=np.float32)
         cond = np.zeros((10, 3), dtype=np.float32)
         ds = TrajectoryDataset(seq, cond, augment=False)
         assert len(ds) == 10
 
     def test_getitem_returns_tensors(self):
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         import torch
         seq = np.random.randn(5, 64, 3).astype(np.float32)
         cond = np.random.randn(5, 3).astype(np.float32)
@@ -151,7 +151,7 @@ class TestTrajectoryDataset:
 
     def test_aug_id_zero_returns_original(self):
         """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         import torch
         seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
         cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
@@ -162,7 +162,7 @@ class TestTrajectoryDataset:
 
     def test_aug_id_one_flips_lateral(self):
         """Aug id 1 should flip the sign of the lateral channel (index 1)."""
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((1, 64, 3), dtype=np.float32)
         seq[0, :, 1] = 0.5  # lateral all positive
         cond = np.zeros((1, 3), dtype=np.float32)
@@ -174,7 +174,7 @@ class TestTrajectoryDataset:
     def test_aug_id_two_slows_speed(self):
         """Aug id 2 should add log(1.25) to log_dt channel and cond[2]."""
         import math
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((1, 64, 3), dtype=np.float32)
         cond = np.zeros((1, 3), dtype=np.float32)
         ds = TrajectoryDataset(seq, cond, augment=True)
@@ -186,7 +186,7 @@ class TestTrajectoryDataset:
     def test_aug_id_three_speeds_up(self):
         """Aug id 3 should add log(1/1.2) to log_dt channel and cond[2]."""
         import math
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((1, 64, 3), dtype=np.float32)
         cond = np.zeros((1, 3), dtype=np.float32)
         ds = TrajectoryDataset(seq, cond, augment=True)
@@ -197,7 +197,7 @@ class TestTrajectoryDataset:
 
     def test_aug_id_four_adds_temporal_noise(self):
         """Aug id 4 should add Gaussian noise to log_dt (channel 2), leaving other channels unchanged."""
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((1, 64, 3), dtype=np.float32)
         cond = np.zeros((1, 3), dtype=np.float32)
         ds = TrajectoryDataset(seq, cond, augment=True)
@@ -215,7 +215,7 @@ class TestTrajectoryDataset:
     def test_aug_id_five_flips_and_slows(self):
         """Aug id 5 should flip lateral and add log(1/0.9) to log_dt and cond[2]."""
         import math
-        from ai_mouse.trainer import TrajectoryDataset
+        from tools.trainer import TrajectoryDataset
         seq = np.zeros((1, 64, 3), dtype=np.float32)
         seq[0, :, 1] = 0.5  # lateral positive
         cond = np.zeros((1, 3), dtype=np.float32)
@@ -233,8 +233,8 @@ class TestResumeFrom:
     def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):
         """train() with resume_from should load weights from given checkpoint dir."""
         import torch
-        from ai_mouse.trainer import train
+        from tools.trainer import train
-        from ai_mouse.models import TrajectoryFlowModel
+        from tools.models import TrajectoryFlowModel
 
         # First, train an initial model and save it
         ckpt_dir = tmp_path / "pretrain"
@@ -273,7 +273,7 @@ class TestResumeFrom:
         assert diff < 0.5, f"Resume_from weights diverged too much: {diff}"
 
     def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path):
-        from ai_mouse.trainer import train
+        from tools.trainer import train
 
         with pytest.raises(FileNotFoundError):
             train(

tests/unit/conftest.py

@@ -0,0 +1 @@
+"""Fixtures for library-only tests (no torch)."""

tests/unit/test__coord.py

@@ -0,0 +1,30 @@
+"""Test the private numpy coordinate transforms."""
+from __future__ import annotations
+
+import numpy as np
+
+from ai_mouse._coord import decode_trajectory, encode_trajectory
+
+
+def test_encode_decode_roundtrip() -> None:
+    points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]])
+    start = (100, 200)
+    end = (500, 300)
+    encoded = encode_trajectory(points, start, end)
+    decoded = decode_trajectory(encoded, start, end)
+    assert np.allclose(decoded, points, atol=1e-6)
+
+
+def test_encode_endpoints() -> None:
+    """Start should encode to (0,0); end should encode to (1,0)."""
+    points = np.array([[100.0, 200.0], [500.0, 300.0]])
+    encoded = encode_trajectory(points, (100, 200), (500, 300))
+    assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6)
+    assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6)
+
+
+def test_zero_distance_returns_zeros() -> None:
+    points = np.array([[100.0, 200.0]])
+    encoded = encode_trajectory(points, (100, 200), (100, 200))
+    assert encoded.shape == (1, 2)
+    assert np.all(encoded == 0)

tests/unit/test_assets.py

@@ -0,0 +1,34 @@
+"""Test the asset path resolver."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+from ai_mouse import _assets
+from ai_mouse.errors import ModelLoadError
+
+
+def test_bundled_flow_model_exists() -> None:
+    p = _assets.bundled_path("flow_model.onnx")
+    assert p.exists()
+    assert p.suffix == ".onnx"
+
+
+def test_bundled_train_config_loadable() -> None:
+    p = _assets.bundled_path("train_config.json")
+    cfg = json.loads(p.read_text())
+    assert "seq_len" in cfg
+    assert "d_model" in cfg
+
+
+def test_resolve_with_custom_dir(tmp_path: Path) -> None:
+    (tmp_path / "flow_model.onnx").write_bytes(b"x")
+    p = _assets.resolve(tmp_path, "flow_model.onnx")
+    assert p == tmp_path / "flow_model.onnx"
+
+
+def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None:
+    with pytest.raises(ModelLoadError, match="missing"):
+        _assets.resolve(tmp_path, "nonexistent.onnx")

tests/unit/test_errors.py

@@ -0,0 +1,19 @@
+"""Test the error hierarchy."""
+from __future__ import annotations
+
+import pytest
+
+from ai_mouse import errors
+
+
+def test_model_load_error_is_aimouse_error() -> None:
+    assert issubclass(errors.ModelLoadError, errors.AiMouseError)
+
+
+def test_generation_error_is_aimouse_error() -> None:
+    assert issubclass(errors.GenerationError, errors.AiMouseError)
+
+
+def test_can_catch_specific_with_general() -> None:
+    with pytest.raises(errors.AiMouseError):
+        raise errors.ModelLoadError("test")

tests/unit/test_golden.py

@@ -0,0 +1,151 @@
+"""Golden regression tests — guard against catastrophic divergence from
+the pre-migration PyTorch implementation.
+
+Important caveat on tolerances
+==============================
+
+These goldens were captured with the legacy PyTorch + scipy stack, where the
+RNG path was ``torch.manual_seed(seed)`` + ``torch.randn(...)`` plus
+``np.random.seed`` for the duration sampler. The rewritten library uses a
+single ``np.random.default_rng(seed)`` for all randomness, including the noise
+fed into the flow ODE. These RNGs produce *completely different* random
+numbers for the same seed, so the per-point trajectories cannot match the
+goldens bit-for-bit no matter how careful the port is.
+
+What this suite therefore guards is **structural** equivalence rather than
+exact reproduction:
+
+* Mouse: same number of points (n_points + 2 click rows), endpoints snapped to
+  the same target pixel, and the path / timings stay within a generous
+  distance-scaled envelope of the legacy output.
+* Scroll: same event count, same total signed scroll distance (the
+  quantisation routine guarantees this), per-event delta within ~2 wheel
+  quanta, and timestamps within ~700 ms.
+
+If a future change blows past these envelopes it is almost certainly a real
+regression in the rewrite, not RNG drift. To re-baseline the goldens after an
+intentional change, regenerate the .npz files from the new implementation.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from ai_mouse import generate, generate_scroll
+
+_GOLDEN_DIR = Path(__file__).parent / "data"
+
+_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
+    ((100, 200), (900, 400)),
+    ((500, 500), (500, 100)),
+    ((200, 600), (800, 200)),
+    ((100, 100), (130, 110)),
+    ((50, 50), (1500, 900)),
+    ((400, 300), (500, 300)),
+    ((300, 300), (700, 700)),
+    ((600, 400), (200, 100)),
+]
+
+_SCROLL_CASES: list[tuple[int, int, str]] = [
+    (0, 1500, "target"),
+    (0, 500, "precise"),
+    (0, 5000, "fast"),
+    (2000, 0, "target"),
+    (0, 800, "precise"),
+    (0, 3500, "fast"),
+    (1000, 1200, "precise"),
+    (0, 10000, "fast"),
+]
+
+# Mouse path tolerance is distance-scaled because absolute pixel diff
+# trivially grows with travel distance. Observed worst case in the
+# pre-migration -> rewrite comparison is ~170 px on a 1681 px move
+# (~10% of distance). 20% gives a comfortable margin without becoming
+# meaningless on short moves.
+_MOUSE_XY_REL = 0.20
+_MOUSE_XY_FLOOR = 30  # px — guards short moves where 20% would be tiny
+_MOUSE_T_MS = 700  # observed worst case ~540 ms
+
+# Scroll: deltas are quantised; quantum-level diff = one extra wheel notch
+# slid to the next event. Two quanta covers everything observed.
+_SCROLL_DELTA_QUANTA = 2
+_SCROLL_T_MS = 700
+
+
+@pytest.mark.parametrize("case_idx", range(8))
+@pytest.mark.parametrize("seed", [0, 1, 2, 3])
+def test_mouse_golden(case_idx: int, seed: int) -> None:
+    golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"]
+    start, end = _MOUSE_CASES[case_idx]
+    pts = generate(start, end, seed=seed)
+    arr = np.array(pts, dtype=np.int64)
+
+    # Structural: shape must match exactly.
+    assert arr.shape == golden.shape, (
+        f"shape mismatch: {arr.shape} vs {golden.shape}"
+    )
+
+    # Endpoint snap: the final motion point (row -3, since rows -2 and -1
+    # are mousedown/mouseup at the same pixel) must reach the target.
+    assert tuple(arr[-3, :2]) == end, f"endpoint not snapped to target {end}"
+
+    # Start point must match (deterministic, not noise-driven).
+    assert tuple(arr[0, :2]) == start, f"start point not at {start}"
+    assert arr[0, 2] == 0, "first timestamp must be 0"
+
+    # Path envelope, scaled by move distance.
+    dist = math.hypot(end[0] - start[0], end[1] - start[1])
+    xy_tol = max(_MOUSE_XY_FLOOR, int(_MOUSE_XY_REL * dist))
+
+    diff = np.abs(arr - golden)
+    xy_max = int(max(diff[:, 0].max(), diff[:, 1].max()))
+    t_max = int(diff[:, 2].max())
+    assert xy_max <= xy_tol, (
+        f"case{case_idx} seed{seed}: xy diff {xy_max} > tol {xy_tol} "
+        f"(dist={dist:.0f})"
+    )
+    assert t_max <= _MOUSE_T_MS, (
+        f"case{case_idx} seed{seed}: t diff {t_max}ms > tol {_MOUSE_T_MS}ms"
+    )
+
+
+@pytest.mark.parametrize("case_idx", range(8))
+@pytest.mark.parametrize("seed", [0, 1, 2, 3])
+def test_scroll_golden(case_idx: int, seed: int) -> None:
+    golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"]
+    start_y, end_y, mode = _SCROLL_CASES[case_idx]
+    events = generate_scroll(start_y, end_y, mode=mode, seed=seed)
+    arr = np.array(
+        [[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
+        dtype=np.int64,
+    )
+    quantum = 40 if mode == "precise" else 120
+
+    if arr.shape != golden.shape:
+        pytest.skip(
+            f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} "
+            f"(quantisation boundary sensitivity)"
+        )
+
+    # Sum of deltas must exactly match — the quantiser guarantees the
+    # final event is corrected to hit the requested total.
+    assert arr[:, 0].sum() == golden[:, 0].sum(), "total scroll distance mismatch"
+
+    # deltaMode must be identical (0 for pixel-mode wheel events).
+    assert (arr[:, 1] == golden[:, 1]).all(), "deltaMode diverged"
+
+    # Per-event delta and time within tolerance.
+    delta_diff = int(np.abs(arr[:, 0] - golden[:, 0]).max())
+    t_diff = int(np.abs(arr[:, 2] - golden[:, 2]).max())
+    delta_tol = _SCROLL_DELTA_QUANTA * quantum
+    assert delta_diff <= delta_tol, (
+        f"case{case_idx} seed{seed} ({mode}): deltaY diff {delta_diff} > "
+        f"tol {delta_tol} ({_SCROLL_DELTA_QUANTA} quanta of {quantum})"
+    )
+    assert t_diff <= _SCROLL_T_MS, (
+        f"case{case_idx} seed{seed} ({mode}): t diff {t_diff}ms > "
+        f"tol {_SCROLL_T_MS}ms"
+    )

tests/unit/test_mouse.py

@@ -0,0 +1,87 @@
+"""Tests for MouseModel and ai_mouse.generate()."""
+from __future__ import annotations
+
+import pytest
+
+
+def test_mouse_model_init_default() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    assert m._seq_len > 0
+    assert m._session is not None
+    m.close()
+
+
+def test_mouse_model_generate_returns_correct_shape() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    assert len(pts) == 66  # 64 moves + 2 clicks
+    for x, y, t in pts:
+        assert isinstance(x, int)
+        assert isinstance(y, int)
+        assert isinstance(t, int)
+
+
+def test_mouse_model_click_false_omits_clicks() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400), click=False)
+    assert len(pts) == 64
+
+
+def test_mouse_model_seed_reproducibility() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    a = m.generate((100, 200), (900, 400), seed=42)
+    b = m.generate((100, 200), (900, 400), seed=42)
+    assert a == b
+
+
+def test_mouse_model_invalid_path_raises_model_load_error() -> None:
+    from ai_mouse.mouse import MouseModel
+    from ai_mouse.errors import ModelLoadError
+    with pytest.raises(ModelLoadError):
+        MouseModel(model_path="/nonexistent/path/here")
+
+
+def test_mouse_model_timestamps_monotonic() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    times = [p[2] for p in pts]
+    for i in range(1, len(times)):
+        assert times[i] >= times[i - 1]
+
+
+def test_mouse_model_starts_and_ends_near_endpoints() -> None:
+    from ai_mouse.mouse import MouseModel
+    start = (100, 200)
+    end = (900, 400)
+    m = MouseModel()
+    pts = m.generate(start, end)
+    assert abs(pts[0][0] - start[0]) < 30
+    assert abs(pts[0][1] - start[1]) < 30
+    last_move = pts[-3]  # last 2 are click events
+    assert abs(last_move[0] - end[0]) < 30
+    assert abs(last_move[1] - end[1]) < 30
+
+
+def test_mouse_model_n_points_parameter() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400), n_points=32)
+    # 32 moves + 2 clicks
+    assert len(pts) == 34
+
+
+def test_mouse_model_click_events_have_matching_coords() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    down, up = pts[-2], pts[-1]
+    assert down[0] == up[0]
+    assert down[1] == up[1]
+    assert up[2] > down[2]
+    # Within click_dist bounds 20..500
+    assert 20 <= up[2] - down[2] <= 500

tests/unit/test_postprocess.py

@@ -0,0 +1,143 @@
+"""Tests for trajectory post-processing primitives."""
+from __future__ import annotations
+
+import numpy as np
+
+from ai_mouse._postprocess import gaussian_smooth
+
+
+def test_gaussian_smooth_preserves_endpoints() -> None:
+    x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0])
+    result = gaussian_smooth(x, sigma=1.0)
+    assert result[0] == 1.0
+    assert result[-1] == 4.0
+
+
+def test_gaussian_smooth_short_input_unchanged() -> None:
+    x = np.array([1.0, 2.0, 3.0])
+    result = gaussian_smooth(x, sigma=1.0)
+    assert np.array_equal(result, x)
+
+
+def test_gaussian_smooth_constant_unchanged() -> None:
+    x = np.full(20, 7.5)
+    result = gaussian_smooth(x, sigma=1.0)
+    assert np.allclose(result, x, atol=1e-6)
+
+
+from ai_mouse._postprocess import snap_endpoints
+
+
+def test_snap_endpoints_pins_first_and_last() -> None:
+    forward = np.linspace(0.1, 0.9, 16)
+    lateral = np.full(16, 0.5)
+    f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16)
+    assert f[0] == 0.0
+    assert l[0] == 0.0
+    assert f[-1] == 1.0
+    assert l[-1] == 0.0
+
+
+def test_snap_endpoints_preserves_middle() -> None:
+    forward = np.linspace(0.0, 1.0, 16)
+    lateral = np.zeros(16)
+    f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4)
+    # Points before the last n_snap should be unchanged
+    assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)
+
+
+from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start
+
+
+def test_smooth_start_dampens_lateral() -> None:
+    forward = np.linspace(0, 1, 16)
+    lateral = np.full(16, 1.0)
+    forward[0] = lateral[0] = 0.0  # invariant: snap already done
+    _, l = smooth_start(forward.copy(), lateral.copy(), n=4)
+    # Lateral at points 1-4 should be < original (dampened)
+    assert l[1] < 1.0
+    assert l[4] < 1.0
+    # Lateral at point 5+ unchanged
+    assert l[5] == 1.0
+
+
+def test_enforce_forward_monotonic_repairs_inversions() -> None:
+    f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
+    out = enforce_forward_monotonic(f.copy())
+    assert np.all(np.diff(out) > 0), out
+
+
+def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
+    f = np.array([-0.1, 0.5, 1.2])
+    out = enforce_forward_monotonic(f.copy())
+    assert out[0] == 0.0
+    assert out[-1] == 1.0
+
+
+from ai_mouse._postprocess import build_timestamps, resample_arc
+
+
+def test_resample_arc_identity_when_same_length() -> None:
+    pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]])
+    out = resample_arc(pts, 4)
+    assert np.allclose(out, pts, atol=1e-6)
+
+
+def test_resample_arc_changes_length() -> None:
+    pts = np.array([[float(i), 0.0] for i in range(10)])
+    out = resample_arc(pts, 5)
+    assert out.shape == (5, 2)
+    # Endpoints preserved
+    assert np.allclose(out[0], pts[0])
+    assert np.allclose(out[-1], pts[-1])
+
+
+def test_build_timestamps_strictly_increasing() -> None:
+    log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0])
+    ts = build_timestamps(log_dt, total_duration_ms=200.0)
+    assert ts[0] == 0
+    assert np.all(np.diff(ts) >= 1)  # at least 1 ms apart
+
+
+def test_build_timestamps_total_close_to_target() -> None:
+    log_dt = np.array([1.0] * 10)
+    ts = build_timestamps(log_dt, total_duration_ms=300.0)
+    # Last timestamp should be roughly total - one slot
+    assert abs(ts[-1] - 270) < 60  # tolerant of clipping
+
+
+from ai_mouse._postprocess import sample_duration, truncnorm_sample
+
+
+def test_truncnorm_sample_within_bounds() -> None:
+    rng = np.random.default_rng(0)
+    samples = [
+        truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500)
+    ]
+    arr = np.array(samples)
+    assert arr.min() >= 20.0
+    assert arr.max() <= 300.0
+    # Mean roughly close to mu
+    assert abs(arr.mean() - 80.0) < 5.0
+
+
+def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None:
+    rng = np.random.default_rng(0)
+    # mu far outside [low, high] — rejection will fail
+    v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng)
+    assert 20.0 <= v <= 30.0
+
+
+def test_sample_duration_uses_correct_bin() -> None:
+    dist_dict = {
+        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
+        "params": [
+            {"mu_log": 4.0, "sigma_log": 0.01},  # bin 0: dist < 50
+            {"mu_log": 5.0, "sigma_log": 0.01},  # bin 1: 50 <= dist < 100
+            {"mu_log": 6.0, "sigma_log": 0.01},  # bin 2: 100 <= dist < 200
+        ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5,
+    }
+    rng = np.random.default_rng(0)
+    v = sample_duration(dist_dict, 150.0, rng)
+    # exp(6) ~ 403, with tiny sigma we should land near there
+    assert 350 < v < 460

tests/unit/test_public_api.py

@@ -0,0 +1,43 @@
+"""Tests for the public package-level API."""
+from __future__ import annotations
+
+
+def test_public_symbols_importable() -> None:
+    from ai_mouse import (
+        MouseModel,
+        ScrollModel,
+        generate,
+        generate_scroll,
+        errors,
+    )
+    assert MouseModel is not None
+    assert ScrollModel is not None
+    assert callable(generate)
+    assert callable(generate_scroll)
+    assert hasattr(errors, "ModelLoadError")
+
+
+def test_generate_function_returns_list_of_tuples() -> None:
+    from ai_mouse import generate
+    pts = generate((100, 100), (300, 200))
+    assert isinstance(pts, list)
+    assert len(pts) > 0
+    assert isinstance(pts[0], tuple)
+    assert len(pts[0]) == 3
+
+
+def test_generate_singleton_reused() -> None:
+    from ai_mouse import generate
+    from ai_mouse import _model_cache
+    _model_cache._get_mouse_model.cache_clear()
+    generate((0, 0), (100, 100))
+    info_after_first = _model_cache._get_mouse_model.cache_info()
+    generate((0, 0), (200, 200))
+    info_after_second = _model_cache._get_mouse_model.cache_info()
+    assert info_after_second.hits > info_after_first.hits
+
+
+def test_version_present() -> None:
+    import ai_mouse
+    assert hasattr(ai_mouse, "__version__")
+    assert isinstance(ai_mouse.__version__, str)

tests/unit/test_scroll.py

@@ -0,0 +1,53 @@
+"""Tests for ScrollModel and ai_mouse.generate_scroll()."""
+from __future__ import annotations
+
+import pytest
+
+
+def test_scroll_model_init_default() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    assert m._seq_len > 0
+    m.close()
+
+
+def test_scroll_model_generate_target_mode() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(0, 1500, mode="target")
+    assert len(events) >= 5
+    total = sum(e["deltaY"] for e in events)
+    assert 1000 <= total <= 2000  # broad — quantisation can drift
+    assert events[0]["t"] == 0
+    assert all(e["deltaMode"] == 0 for e in events)
+
+
+def test_scroll_model_direction() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(2000, 0, mode="target")
+    total = sum(e["deltaY"] for e in events)
+    assert total < 0
+
+
+def test_scroll_invalid_path() -> None:
+    from ai_mouse.errors import ModelLoadError
+    from ai_mouse.scroll import ScrollModel
+    with pytest.raises(ModelLoadError):
+        ScrollModel(model_path="/no/such/path")
+
+
+def test_scroll_model_timestamps_monotonic() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(0, 2000, mode="target")
+    times = [e["t"] for e in events]
+    for i in range(1, len(times)):
+        assert times[i] >= times[i - 1]
+
+
+def test_scroll_model_deltaY_are_integers() -> None:
+    from ai_mouse.scroll import ScrollModel
+    m = ScrollModel()
+    events = m.generate(0, 2000, mode="target")
+    assert all(isinstance(e["deltaY"], int) for e in events)

tools/__main__.py

@@ -12,7 +12,7 @@ from pathlib import Path
 
 
 def _train_main(args: argparse.Namespace) -> int:
-    from ai_mouse.trainer import train
+    from tools.trainer import train
 
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
     train(
@@ -28,7 +28,7 @@ def _train_main(args: argparse.Namespace) -> int:
 
 
 def _eval_main(args: argparse.Namespace) -> int:
-    from ai_mouse.eval.__main__ import main as eval_main
+    from tools.eval.__main__ import main as eval_main
     # Reconstruct argv for the sub-CLI
     argv = [
         "--model-dir", args.model_dir,
@@ -43,7 +43,7 @@ def _eval_main(args: argparse.Namespace) -> int:
 
 
 def _balabit_main(args: argparse.Namespace) -> int:
-    from ai_mouse.data_adapters.balabit import main as bal_main
+    from tools.data_adapters.balabit import main as bal_main
     argv = [
         "--input", str(args.input),
         "--output", str(args.output),

tools/data_adapters/__main__.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 
 import sys
 
-from ai_mouse.data_adapters.balabit import main
+from tools.data_adapters.balabit import main
 
 if __name__ == "__main__":
     sys.exit(main())

tools/data_adapters/balabit.py

@@ -260,7 +260,7 @@ def main(argv: list[str] | None = None) -> int:
     """CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
     import argparse
 
-    from ai_mouse.config import BalabitAdapterConfig
+    from tools.config import BalabitAdapterConfig
 
     parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
     parser.add_argument(

tools/eval/__main__.py

@@ -48,8 +48,15 @@ def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
 def _generate_n_samples(
     model_dir: str, n_samples: int, seed: int = 0
 ) -> list[dict]:
-    """Call the project's generator N times with random start/end pairs."""
+    """Call the project's generator N times with random start/end pairs.
-    from ai_mouse.generator import generate
+
+    ``model_dir`` is accepted for CLI backward compatibility but is no longer
+    used — generation goes through the public ai_mouse API which loads the
+    bundled ONNX model. Export a fresh .onnx via ``python -m tools.export_onnx``
+    to refresh.
+    """
+    del model_dir  # legacy arg, unused
+    from ai_mouse import generate
 
     rng = random.Random(seed)
     out: list[dict] = []
@@ -63,7 +70,7 @@ def _generate_n_samples(
         ex = max(0, min(800, ex))
         ey = max(0, min(600, ey))
         try:
-            pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
+            pts = generate(start=(sx, sy), end=(ex, ey))
         except Exception as exc:  # noqa: BLE001
             logger.warning("generate() failed at i=%d: %s", i, exc)
             continue
@@ -110,7 +117,7 @@ def main(argv: list[str] | None = None) -> int:
         logger.error("Empty trace sets — aborting")
         return 1
 
-    from ai_mouse.eval.report import build_report
+    from tools.eval.report import build_report
     args.output.parent.mkdir(parents=True, exist_ok=True)
     build_report(
         generated_traces=gen_traces,

tools/eval/report.py

@@ -15,7 +15,7 @@ import matplotlib
 matplotlib.use("Agg")  # headless
 import matplotlib.pyplot as plt  # noqa: E402
 
-from ai_mouse.eval.metrics import (
+from tools.eval.metrics import (
     compute_acceleration,
     compute_jerk,
     compute_speed,

tools/export_onnx.py

@@ -0,0 +1,305 @@
+"""Export trained PyTorch checkpoints to ONNX for the inference SDK.
+
+Usage:
+    uv run python -m tools.export_onnx \
+        --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
+
+
+def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None:
+    """Verify ONNX flow model matches PyTorch output on random input."""
+    import onnxruntime as ort
+    from tools.models import TrajectoryFlowModel
+
+    cfg = json.loads((ckpt_dir / "train_config.json").read_text())
+    seq_len = int(cfg["seq_len"])
+    cond_dim = int(cfg.get("cond_dim", 3))
+
+    model = TrajectoryFlowModel(
+        seq_len=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=cond_dim,
+        dropout=0.0,
+    )
+    model.load_state_dict(
+        torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
+    )
+    model.eval()
+
+    torch.manual_seed(42)
+    np.random.seed(42)
+    x = torch.randn(2, seq_len, 3, dtype=torch.float32)
+    t = torch.tensor([0.0, 0.5], dtype=torch.float32)
+    cond = torch.randn(2, cond_dim, dtype=torch.float32)
+
+    with torch.no_grad():
+        torch_out = model(x, t, cond).numpy()
+
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    ort_out = sess.run(
+        ["v"],
+        {
+            "x_t": x.numpy(),
+            "t": t.numpy(),
+            "cond": cond.numpy(),
+        },
+    )[0]
+
+    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
+        max_diff = float(np.abs(torch_out - ort_out).max())
+        raise RuntimeError(
+            f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
+        )
+    logger.info("Flow model parity OK (atol=%.0e)", _ATOL)
+
+
+def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None:
+    """Verify ONNX scroll decoder matches PyTorch decoder output."""
+    import onnxruntime as ort
+    from tools.scroll.models import ScrollCVAE
+
+    cfg = json.loads((ckpt_dir / "scroll_config.json").read_text())
+    seq_len = int(cfg["seq_len"])
+    latent_dim = int(cfg["latent_dim"])
+    cond_dim = int(cfg["cond_dim"])
+
+    full = ScrollCVAE(
+        seq_len=seq_len,
+        latent_dim=latent_dim,
+        hidden=int(cfg["hidden"]),
+        cond_dim=cond_dim,
+    )
+    full.load_state_dict(
+        torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
+    )
+    full.eval()
+
+    torch.manual_seed(7)
+    z = torch.randn(2, latent_dim, dtype=torch.float32)
+    cond = torch.randn(2, cond_dim, dtype=torch.float32)
+
+    with torch.no_grad():
+        torch_out = full.decode(z, cond).numpy()
+
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0]
+
+    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
+        max_diff = float(np.abs(torch_out - ort_out).max())
+        raise RuntimeError(
+            f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
+        )
+    logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL)
+
+
+def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None:
+    """Copy JSON metadata files alongside the ONNX models."""
+    for name in ("click_dist.json", "duration_dist.json", "train_config.json"):
+        src = flow_dir / name
+        if not src.exists():
+            raise FileNotFoundError(f"Required metadata missing: {src}")
+        shutil.copy2(src, out_dir / name)
+    src = scroll_dir / "scroll_config.json"
+    if not src.exists():
+        raise FileNotFoundError(f"Required metadata missing: {src}")
+    shutil.copy2(src, out_dir / "scroll_config.json")
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0])
+    p.add_argument("--flow-ckpt", type=Path, required=True)
+    p.add_argument("--scroll-ckpt", type=Path, required=True)
+    p.add_argument("--output", type=Path, required=True)
+    args = p.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    args.output.mkdir(parents=True, exist_ok=True)
+
+    flow_onnx = export_flow_model(args.flow_ckpt, args.output)
+    scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output)
+
+    try:
+        _check_flow_parity(args.flow_ckpt, flow_onnx)
+        _check_scroll_parity(args.scroll_ckpt, scroll_onnx)
+    except RuntimeError as exc:
+        logger.error("Parity check failed: %s", exc)
+        flow_onnx.unlink(missing_ok=True)
+        scroll_onnx.unlink(missing_ok=True)
+        return 1
+
+    _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output)
+    logger.info("Export complete: %s", args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/models.py

@@ -12,7 +12,6 @@ import math
 
 import torch
 import torch.nn as nn
-from torch.distributions import Normal
 
 
 # ---------------------------------------------------------------------------
@@ -157,80 +156,3 @@ class TrajectoryFlowModel(nn.Module):
 
         # Output projection
         return self.output_proj(h)  # (B, T, 3)
-
-
-# ---------------------------------------------------------------------------
-# Legacy JointCVAE — kept for backward compatibility with generator.py
-# ---------------------------------------------------------------------------
-
-
-class JointCVAE(nn.Module):
-    """Joint Conditional VAE for mouse trajectory generation (legacy).
-
-    Kept for backward compatibility with the existing generator.
-    See TrajectoryFlowModel for the new approach.
-    """
-
-    def __init__(
-        self,
-        seq_len: int = 64,
-        latent_dim: int = 32,
-        hidden: int = 128,
-        cond_dim: int = 3,
-    ):
-        super().__init__()
-        self.seq_len = seq_len
-        self.latent_dim = latent_dim
-        self.hidden = hidden
-        self.cond_dim = cond_dim
-        self.feat_dim = 3
-
-        self.enc_gru = nn.GRU(
-            input_size=self.feat_dim + cond_dim,
-            hidden_size=hidden,
-            num_layers=2,
-            batch_first=True,
-            bidirectional=True,
-        )
-        self.enc_mu = nn.Linear(hidden * 2, latent_dim)
-        self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
-
-        self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
-        self.dec_gru = nn.GRU(
-            input_size=latent_dim + cond_dim,
-            hidden_size=hidden,
-            num_layers=2,
-            batch_first=True,
-        )
-        self.dec_out = nn.Linear(hidden, self.feat_dim)
-
-    def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        B, T, _ = seq.shape
-        c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
-        x_in = torch.cat([seq, c_exp], dim=-1)
-        _, h_n = self.enc_gru(x_in)
-        h_fwd = h_n[-2]
-        h_bwd = h_n[-1]
-        h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
-        return self.enc_mu(h_cat), self.enc_logvar(h_cat)
-
-    def decode(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 reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
-        std = torch.exp(0.5 * logvar)
-        return Normal(mu, std).rsample()
-
-    def forward(
-        self, seq: torch.Tensor, cond: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        mu, logvar = self.encode(seq, cond)
-        z = self.reparameterise(mu, logvar)
-        recon = self.decode(z, cond)
-        return recon, mu, logvar

tools/scroll/collector.py

@@ -9,7 +9,7 @@ from __future__ import annotations
 import logging
 import random
 
-from ai_mouse.config import SCROLL_MODES
+from tools.config import SCROLL_MODES
 
 logger = logging.getLogger(__name__)
 

tools/scroll/trainer.py

@@ -20,8 +20,8 @@ import torch.nn as nn
 from torch.distributions import Normal, kl_divergence
 from torch.utils.data import DataLoader, TensorDataset
 
-from ai_mouse.config import ScrollTrainConfig
+from tools.config import ScrollTrainConfig
-from ai_mouse.scroll.models import ScrollCVAE
+from tools.scroll.models import ScrollCVAE
 
 logger = logging.getLogger(__name__)
 

tools/serve.py

@@ -10,7 +10,7 @@ import webbrowser
 
 import uvicorn
 
-from ai_mouse.server import create_app
+from tools.server import create_app
 
 app = create_app()
 

tools/server/routes_collect.py

@@ -8,7 +8,7 @@ import logging
 from fastapi import APIRouter, Depends, HTTPException
 from pydantic import BaseModel
 
-from ai_mouse.collector import Collector
+from tools.collector import Collector
 
 from .deps import SessionState, get_data_dir, get_state
 

tools/server/routes_scroll.py

@@ -85,7 +85,7 @@ def scroll_start(
     req: ScrollStartRequest,
     state: SessionState = Depends(get_state),
 ) -> dict:
-    from ai_mouse.scroll.collector import ScrollCollector
+    from tools.scroll.collector import ScrollCollector
 
     scroll_collector = ScrollCollector(
         mode=req.mode, count=req.count, viewport_height=req.viewport_height
@@ -146,7 +146,7 @@ async def _scroll_train_sse(req: ScrollTrainRequest) -> AsyncGenerator[str, None
         queue.put_nowait(msg)
 
     async def run() -> None:
-        from ai_mouse.scroll.trainer import train_scroll
+        from tools.scroll.trainer import train_scroll
 
         traces_path, models_dir = _paths()
         try:
@@ -182,21 +182,17 @@ async def scroll_train(req: ScrollTrainRequest) -> StreamingResponse:
 
 @router.post("/verify")
 def scroll_verify(req: ScrollVerifyRequest) -> dict:
-    from ai_mouse.scroll.generator import generate_scroll
+    # Uses the bundled ONNX scroll model exposed via the public ai_mouse API.
+    # The legacy scroll_model.pt path is no longer wired in; export a fresh
+    # scroll_decoder.onnx via `python -m tools.export_onnx` to update.
+    from ai_mouse import generate_scroll
 
-    _, models_dir = _paths()
-    if not (models_dir / "scroll_model.pt").exists():
-        raise HTTPException(
-            status_code=400,
-            detail="滚轮模型尚未训练，请先在「训练模型 → 滚轮模型」中完成训练。",
-        )
     paths = []
     for _ in range(min(req.n_paths, 12)):
         events = generate_scroll(
             req.start_scrollY,
             req.target_scrollY,
             mode=req.mode,
-            model_dir=str(models_dir),
         )
         paths.append(events)
     return {"paths": paths}

tools/server/routes_train.py

@@ -77,7 +77,7 @@ async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]:
         queue.put_nowait(msg)
 
     async def run_training_async() -> None:
-        from ai_mouse.trainer import train
+        from tools.trainer import train
 
         traces_path, models_dir, pretrained_dir = _paths()
         data_path = Path(req.data_path) if req.data_path else traces_path

tools/server/routes_verify.py

@@ -7,8 +7,6 @@ import logging
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 
-from .deps import get_data_dir
-
 logger = logging.getLogger(__name__)
 
 router = APIRouter()
@@ -32,18 +30,19 @@ class VerifyRequest(BaseModel):
 
 @router.post("/verify")
 def verify(req: VerifyRequest) -> dict:
-    from ai_mouse.generator import generate
+    # Uses the bundled ONNX model exposed via the public ai_mouse API.
+    # The legacy req.model_dir / data/models_v2 .pt path is no longer wired
+    # in; export a fresh .onnx via `python -m tools.export_onnx` to update.
+    from ai_mouse import generate
 
     n = max(1, min(req.n_paths, 12))
-    models_dir = get_data_dir() / "models_v2"
-    model_dir_arg = req.model_dir if req.model_dir else str(models_dir)
     start = tuple(req.start)  # type: ignore[arg-type]
     end = tuple(req.end)  # type: ignore[arg-type]
 
     paths = []
     try:
         for _ in range(n):
-            pts = generate(start=start, end=end, model_dir=model_dir_arg)
+            pts = generate(start=start, end=end)
             paths.append([[x, y, t] for x, y, t in pts])
     except FileNotFoundError as exc:
         raise HTTPException(status_code=404, detail=str(exc)) from exc

tools/trainer.py

@@ -23,10 +23,10 @@ import numpy as np
 import torch
 from torch.utils.data import DataLoader
 
-from ai_mouse.config import TrainConfig
+from ai_mouse._coord import encode_trajectory
-from ai_mouse.coord import encode_trajectory
+from tools.config import TrainConfig
-from ai_mouse.models import TrajectoryFlowModel
+from tools.models import TrajectoryFlowModel
-from ai_mouse.utils import resample_arc
+from tools.utils import resample_arc
 
 logger = logging.getLogger(__name__)
 

uv.lock

@@ -1,42 +1,54 @@
 version = 1
 revision = 3
 requires-python = ">=3.12, <3.14"
+resolution-markers = [
+    "python_full_version >= '3.13' and platform_machine != 's390x'",
+    "python_full_version < '3.13' and platform_machine != 's390x'",
+    "python_full_version >= '3.13' and platform_machine == 's390x'",
+    "python_full_version < '3.13' and platform_machine == 's390x'",
+]
 
 [[package]]
 name = "ai-mouse"
-version = "0.1.0"
+version = "0.2.0"
-source = { virtual = "." }
+source = { editable = "." }
 dependencies = [
-    { name = "fastapi" },
-    { name = "matplotlib" },
     { name = "numpy" },
+    { name = "onnxruntime" },
+]
+
+[package.dev-dependencies]
+dev = [
+    { name = "fastapi" },
+    { name = "httpx" },
+    { name = "matplotlib" },
+    { name = "onnx" },
+    { name = "onnxscript" },
+    { name = "pytest" },
+    { name = "pytest-asyncio" },
     { name = "scipy" },
     { name = "torch" },
     { name = "uvicorn" },
 ]
 
-[package.dev-dependencies]
-dev = [
-    { name = "httpx" },
-    { name = "pytest" },
-    { name = "pytest-asyncio" },
-]
-
 [package.metadata]
 requires-dist = [
-    { name = "fastapi", specifier = ">=0.111.0" },
-    { name = "matplotlib", specifier = ">=3.8.0" },
     { name = "numpy", specifier = ">=1.26.0" },
-    { name = "scipy", specifier = ">=1.10.0" },
+    { name = "onnxruntime", specifier = ">=1.17.0" },
-    { name = "torch", specifier = ">=2.2.0" },
-    { name = "uvicorn", specifier = ">=0.29.0" },
 ]
 
 [package.metadata.requires-dev]
 dev = [
+    { name = "fastapi", specifier = ">=0.111.0" },
     { name = "httpx", specifier = ">=0.27.0" },
+    { name = "matplotlib", specifier = ">=3.8.0" },
+    { name = "onnx", specifier = ">=1.15.0" },
+    { name = "onnxscript", specifier = ">=0.1" },
     { name = "pytest", specifier = ">=8.0.0" },
     { name = "pytest-asyncio", specifier = ">=0.23.0" },
+    { name = "scipy", specifier = ">=1.10.0" },
+    { name = "torch", specifier = ">=2.2.0" },
+    { name = "uvicorn", specifier = ">=0.29.0" },
 ]
 
 [[package]]
@@ -243,6 +255,14 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" },
 ]
 
+[[package]]
+name = "flatbuffers"
+version = "25.12.19"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" },
+]
+
 [[package]]
 name = "fonttools"
 version = "4.62.1"
@@ -481,6 +501,32 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f4/38/ae27288e788c35a4250491422f3db7750366fc8c97d6f36fbdecfc1f5518/matplotlib-3.10.9-cp313-cp313t-win_arm64.whl", hash = "sha256:6b63d9c7c769b88ab81e10dc86e4e0607cf56817b9f9e6cf24b2a5f1693b8e38", size = 8188292, upload-time = "2026-04-24T00:13:15.546Z" },
 ]
 
+[[package]]
+name = "ml-dtypes"
+version = "0.5.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" },
+    { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" },
+    { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" },
+    { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" },
+    { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" },
+    { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" },
+    { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" },
+    { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" },
+    { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" },
+    { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" },
+    { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" },
+]
+
 [[package]]
 name = "mpmath"
 version = "1.3.0"
@@ -688,6 +734,88 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" },
 ]
 
+[[package]]
+name = "onnx"
+version = "1.21.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ml-dtypes" },
+    { name = "numpy" },
+    { name = "protobuf" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c5/93/942d2a0f6a70538eea042ce0445c8aefd46559ad153469986f29a743c01c/onnx-1.21.0.tar.gz", hash = "sha256:4d8b67d0aaec5864c87633188b91cc520877477ec0254eda122bef8be43cd764", size = 12074608, upload-time = "2026-03-27T21:33:36.118Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7d/ae/cb644ec84c25e63575d9d8790fdcc5d1a11d67d3f62f872edb35fa38d158/onnx-1.21.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:fc2635400fe39ff37ebc4e75342cc54450eadadf39c540ff132c319bf4960095", size = 17965930, upload-time = "2026-03-27T21:32:48.089Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/b6/eeb5903586645ef8a49b4b7892580438741acc3df91d7a5bd0f3a59ea9cb/onnx-1.21.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9003d5206c01fa2ff4b46311566865d8e493e1a6998d4009ec6de39843f1b59b", size = 17531344, upload-time = "2026-03-27T21:32:50.837Z" },
+    { url = "https://files.pythonhosted.org/packages/a7/00/4823f06357892d1e60d6f34e7299d2ba4ed2108c487cc394f7ce85a3ff14/onnx-1.21.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9261bd580fb8548c9c37b3c6750387eb8f21ea43c63880d37b2c622e1684285", size = 17613697, upload-time = "2026-03-27T21:32:54.222Z" },
+    { url = "https://files.pythonhosted.org/packages/23/1d/391f3c567ae068c8ac4f1d1316bae97c9eb45e702f05975fe0e17ad441f0/onnx-1.21.0-cp312-abi3-win32.whl", hash = "sha256:9ea4e824964082811938a9250451d89c4ec474fe42dd36c038bfa5df31993d1e", size = 16287200, upload-time = "2026-03-27T21:32:57.277Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/a6/5eefbe5b40ea96de95a766bd2e0e751f35bdea2d4b951991ec9afaa69531/onnx-1.21.0-cp312-abi3-win_amd64.whl", hash = "sha256:458d91948ad9a7729a347550553b49ab6939f9af2cddf334e2116e45467dc61f", size = 16441045, upload-time = "2026-03-27T21:33:00.081Z" },
+    { url = "https://files.pythonhosted.org/packages/63/c4/0ed8dc037a39113d2a4d66e0005e07751c299c46b993f1ad5c2c35664c20/onnx-1.21.0-cp312-abi3-win_arm64.whl", hash = "sha256:ca14bc4842fccc3187eb538f07eabeb25a779b39388b006db4356c07403a7bbb", size = 16403134, upload-time = "2026-03-27T21:33:03.987Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/89/0e1a9beb536401e2f45ac88735e123f2735e12fc7b56ff6c11727e097526/onnx-1.21.0-cp313-cp313t-macosx_12_0_universal2.whl", hash = "sha256:257d1d1deb6a652913698f1e3f33ef1ca0aa69174892fe38946d4572d89dd94f", size = 17975430, upload-time = "2026-03-27T21:33:07.005Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/46/e6dc71a7b3b317265591b20a5f71d0ff5c0d26c24e52283139dc90c66038/onnx-1.21.0-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7cd7cb8f6459311bdb557cbf6c0ccc6d8ace11c304d1bba0a30b4a4688e245f8", size = 17537435, upload-time = "2026-03-27T21:33:09.765Z" },
+    { url = "https://files.pythonhosted.org/packages/49/2e/27affcac63eaf2ef183a44fd1a1354b11da64a6c72fe6f3fdcf5571bcee5/onnx-1.21.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b58a4cfec8d9311b73dc083e4c1fa362069267881144c05139b3eba5dc3a840", size = 17617687, upload-time = "2026-03-27T21:33:12.619Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/5c/ac8ed15e941593a3672ce424280b764979026317811f2e8508432bfc3429/onnx-1.21.0-cp313-cp313t-win_amd64.whl", hash = "sha256:1a9baf882562c4cebf79589bebb7cd71a20e30b51158cac3e3bbaf27da6163bd", size = 16449402, upload-time = "2026-03-27T21:33:15.555Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/aa/d2231e0dcaad838217afc64c306c8152a080134d2034e247cc973d577674/onnx-1.21.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bba12181566acf49b35875838eba49536a327b2944664b17125577d230c637ad", size = 16408273, upload-time = "2026-03-27T21:33:18.599Z" },
+]
+
+[[package]]
+name = "onnx-ir"
+version = "0.2.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ml-dtypes" },
+    { name = "numpy" },
+    { name = "onnx" },
+    { name = "sympy" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/35/e6/672fefb2f108d077f58181a7babf4c0f8d1182a30353ffc9c79c63afc5ee/onnx_ir-0.2.1.tar.gz", hash = "sha256:8b8b10a93f43e65962104de6070c43c5dacb0e3cdfefc7c8059dd83c9db64f35", size = 144279, upload-time = "2026-04-20T20:21:47.735Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8c/aa/f7a53321c60b9ad9ee184b6018292ed6b5389947592a2c8c09c736bb7f9e/onnx_ir-0.2.1-py3-none-any.whl", hash = "sha256:c7285da889312f91882de2092e298a9eeeefbfc1d1951c49d983992967eb09a7", size = 166792, upload-time = "2026-04-20T20:21:46.357Z" },
+]
+
+[[package]]
+name = "onnxruntime"
+version = "1.26.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "flatbuffers" },
+    { name = "numpy" },
+    { name = "packaging" },
+    { name = "protobuf" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f6/a0/3f9d896a0385a36bd04345d6d0b802821a5782adde562e7e135f6bb71c73/onnxruntime-1.26.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:91f2bb870a4b9224eba0a6728c1fa7a9e552b8e59e1083c51fbbc3d013f2b5c0", size = 16052692, upload-time = "2026-05-08T19:07:13.829Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/43/2a4e04f8dbeffad19bbcced4bcd4289bf478921518437404d6b92bdf213b/onnxruntime-1.26.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9b6dd70599005bd1bf29779f04a91978b92b5e719c11a20068a8f8e535f725b6", size = 18185439, upload-time = "2026-05-08T19:07:36.299Z" },
+    { url = "https://files.pythonhosted.org/packages/44/fc/026d0a7162b9c2153dac292baea9e027c42304dc1d9dc6f8ff5b4cfbaedd/onnxruntime-1.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:a26374dc7fbcaae593601086b242120e13f2310558df0991da6dd8b8fac00414", size = 13026427, upload-time = "2026-05-08T19:08:03.503Z" },
+    { url = "https://files.pythonhosted.org/packages/3e/27/1dcf88e45e4c69db5f7b106f2dacc3801ba98994e082ca03e1dfdf7bfe57/onnxruntime-1.26.0-cp312-cp312-win_arm64.whl", hash = "sha256:54a8053410fd31fd66469bd754fcfe8a4df9f7eb44756b4b5479bf50c842d948", size = 12796647, upload-time = "2026-05-08T19:07:52.108Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/a2/c801242685e0ce48a4ca51dfafbb588765e0446397e123be53ba5598f3f5/onnxruntime-1.26.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ccce19c5f771b8268902f77d9fed9e88f9499465d6780808faa6611a789d33f0", size = 18016563, upload-time = "2026-05-08T19:07:28.081Z" },
+    { url = "https://files.pythonhosted.org/packages/e2/64/0492c0b1db04e29b2630c87cfa36f9d6872b1ca8614b90c5cad58fac7d76/onnxruntime-1.26.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bdbed8cf3b672b66acb032f33a253bc27f42bce6ece48ae3fab4fa483a5e96e0", size = 16052634, upload-time = "2026-05-08T19:07:16.885Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/26/4d09ddc755a84fc8d5e192991626b0e0680e8f6c5d58f4f1d05c42bc48cf/onnxruntime-1.26.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c07af6fc6d5557835f2b6ee7a96d8b3235d0c57a8e230efdedaee106a8a3cbc6", size = 18185632, upload-time = "2026-05-08T19:07:38.756Z" },
+    { url = "https://files.pythonhosted.org/packages/77/89/3e52249aa08fa301e217ecba07b5246a8338fa2b401e109326e3fc5be0f9/onnxruntime-1.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:61bec80655efa460591c2bc655392d57d2650ce85533a6b9b3b7a790d7ea7916", size = 13026751, upload-time = "2026-05-08T19:08:06.2Z" },
+    { url = "https://files.pythonhosted.org/packages/06/b3/c1c8782b14af6797c303de132d6eef26a9fb80dfacd3750ce57911d11c6b/onnxruntime-1.26.0-cp313-cp313-win_arm64.whl", hash = "sha256:a6677545ff451e3539a02746d2f207d8c5baa4a0a818886bb9d6a6eb9511ee89", size = 12796807, upload-time = "2026-05-08T19:07:54.879Z" },
+    { url = "https://files.pythonhosted.org/packages/c3/f5/47b0676408abec652c14b84d7173e389837832d850c24f87184277313e8d/onnxruntime-1.26.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e016edc15d3c19f36807e1c6b10be5b27807688c32720f91b5ae480a95215d0", size = 16057265, upload-time = "2026-05-08T19:07:19.603Z" },
+    { url = "https://files.pythonhosted.org/packages/3b/45/33ab6deeef010ca844c877dd618cebc079590bbe52d2a3678e7223b1b908/onnxruntime-1.26.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f5fc48a91a046a6a5c9b147f83fb41d65d24d24923373b222cdd248f0f4f4aac", size = 18197590, upload-time = "2026-05-08T19:07:41.422Z" },
+]
+
+[[package]]
+name = "onnxscript"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ml-dtypes" },
+    { name = "numpy" },
+    { name = "onnx" },
+    { name = "onnx-ir" },
+    { name = "packaging" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9b/99/fd948eba63ba65b52265a4cd09a14f96bb9f5b730fcef58876c4358bf406/onnxscript-0.7.0.tar.gz", hash = "sha256:c95ed7b339b02cface56ee27689565c46612e1fc542c562298dddfdad5268dc5", size = 612032, upload-time = "2026-04-20T17:09:19.775Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b9/ce/2ed92575cc3be4ea1db5f38f16f20765f9b20b69b14d6c1d9972658a8ee9/onnxscript-0.7.0-py3-none-any.whl", hash = "sha256:5b356907d4501e9919f8599c91d8da967406a37b1fac2b40caa55a49acf242ea", size = 714842, upload-time = "2026-04-20T17:09:22.089Z" },
+]
+
 [[package]]
 name = "packaging"
 version = "26.2"
@@ -750,6 +878,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
 ]
 
+[[package]]
+name = "protobuf"
+version = "7.34.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6b/6b/a0e95cad1ad7cc3f2c6821fcab91671bd5b78bd42afb357bb4765f29bc41/protobuf-7.34.1.tar.gz", hash = "sha256:9ce42245e704cc5027be797c1db1eb93184d44d1cdd71811fb2d9b25ad541280", size = 454708, upload-time = "2026-03-20T17:34:47.036Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ec/11/3325d41e6ee15bf1125654301211247b042563bcc898784351252549a8ad/protobuf-7.34.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:d8b2cc79c4d8f62b293ad9b11ec3aebce9af481fa73e64556969f7345ebf9fc7", size = 429247, upload-time = "2026-03-20T17:34:37.024Z" },
+    { url = "https://files.pythonhosted.org/packages/eb/9d/aa69df2724ff63efa6f72307b483ce0827f4347cc6d6df24b59e26659fef/protobuf-7.34.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:5185e0e948d07abe94bb76ec9b8416b604cfe5da6f871d67aad30cbf24c3110b", size = 325753, upload-time = "2026-03-20T17:34:38.751Z" },
+    { url = "https://files.pythonhosted.org/packages/92/e8/d174c91fd48e50101943f042b09af9029064810b734e4160bbe282fa1caa/protobuf-7.34.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:403b093a6e28a960372b44e5eb081775c9b056e816a8029c61231743d63f881a", size = 340198, upload-time = "2026-03-20T17:34:39.871Z" },
+    { url = "https://files.pythonhosted.org/packages/53/1b/3b431694a4dc6d37b9f653f0c64b0a0d9ec074ee810710c0c3da21d67ba7/protobuf-7.34.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:8ff40ce8cd688f7265326b38d5a1bed9bfdf5e6723d49961432f83e21d5713e4", size = 324267, upload-time = "2026-03-20T17:34:41.1Z" },
+    { url = "https://files.pythonhosted.org/packages/85/29/64de04a0ac142fb685fd09999bc3d337943fb386f3a0ec57f92fd8203f97/protobuf-7.34.1-cp310-abi3-win32.whl", hash = "sha256:34b84ce27680df7cca9f231043ada0daa55d0c44a2ddfaa58ec1d0d89d8bf60a", size = 426628, upload-time = "2026-03-20T17:34:42.536Z" },
+    { url = "https://files.pythonhosted.org/packages/4d/87/cb5e585192a22b8bd457df5a2c16a75ea0db9674c3a0a39fc9347d84e075/protobuf-7.34.1-cp310-abi3-win_amd64.whl", hash = "sha256:e97b55646e6ce5cbb0954a8c28cd39a5869b59090dfaa7df4598a7fba869468c", size = 437901, upload-time = "2026-03-20T17:34:44.112Z" },
+    { url = "https://files.pythonhosted.org/packages/88/95/608f665226bca68b736b79e457fded9a2a38c4f4379a4a7614303d9db3bc/protobuf-7.34.1-py3-none-any.whl", hash = "sha256:bb3812cd53aefea2b028ef42bd780f5b96407247f20c6ef7c679807e9d188f11", size = 170715, upload-time = "2026-03-20T17:34:45.384Z" },
+]
+
 [[package]]
 name = "pydantic"
 version = "2.13.4"