ai_mouse/docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md

ai_mouse 库化与 ONNX Runtime 迁移 — 设计文档

日期: 2026-05-11 状态: Draft (待用户审阅) 目标版本: ai_mouse 0.2.0

背景与目标

当前 ai_mouse 仓库既是应用又是包:推理代码、训练代码、Web UI、采集、评估、数据适配器都混在 ai_mouse/ 包目录下,运行时强依赖 torch、fastapi、scipy、matplotlib。这与"被其他项目直接引用的库"定位不符——任何只想调用 generate() 的下游必须装一整套 GB 级依赖。

本次重构的目标:

1. 物理边界: 推理库代码与开发期代码完全分开,推理库不依赖 torch
2. ONNX 化: 推理走 ONNX Runtime,预训练权重以 ONNX 格式打进 wheel
3. 可消费: 下游 pip install git+... 即可 from ai_mouse import generate
4. API 优化: 同时提供函数式入口(向后兼容)和类入口(性能场景),底层共享 ORT session

非目标(显式排除):

• 不发布到 PyPI(本阶段仅 git URL 安装)
• 不重写训练代码以脱离 torch
• 不做 ONNX Runtime Training
• 不为视觉真实性设硬阈值断言

约束(由用户决策确认)

维度	决策
使用场景	纯推理 SDK,零 torch 运行时依赖
权重交付	打进 wheel(importlib.resources)
开发期代码位置	仓库顶层 tools/,不进 wheel
公开 API 形态	类 + 函数双轨;函数内部走进程级 lru_cache 单例
发布渠道	git URL 直装,不发 PyPI

设计

§1 公开 API

顶层入口 (src/ai_mouse/__init__.py)

from ai_mouse import (
    MouseModel,       # 类: 常驻 InferenceSession
    ScrollModel,
    generate,         # 便利函数: 进程级单例缓存
    generate_scroll,
)
from ai_mouse import errors  # ModelLoadError, GenerationError
__version__ = "0.2.0"

MouseModel

class MouseModel:
    def __init__(
        self,
        model_path: str | Path | None = None,    # None ⇒ wheel 内置权重
        providers: Sequence[str] | None = None,  # None ⇒ ORT 默认顺序 (CPU)
        seed: int | None = None,                 # None ⇒ 每次调用随机
    ) -> None: ...

    def generate(
        self,
        start: tuple[int, int],
        end: tuple[int, int],
        n_points: int = 64,
        speed: float | None = None,
        click: bool = True,                      # True ⇒ 末尾追加 mouse-down/-up
        seed: int | None = None,                 # 覆盖实例 seed
    ) -> list[tuple[int, int, int]]: ...

    def sample_click_duration_ms(self, seed: int | None = None) -> int: ...

    def close(self) -> None: ...
    def __enter__(self) -> "MouseModel": ...
    def __exit__(self, *exc) -> None: ...

ScrollModel(对称)

class ScrollModel:
    def __init__(self, model_path=None, providers=None, seed=None): ...

    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]: ...

    def close(self) -> None: ...
    def __enter__(self) -> "ScrollModel": ...
    def __exit__(self, *exc) -> None: ...

便利函数

def generate(start, end, *, n_points=64, speed=None, click=True, seed=None,
             model_path=None, providers=None) -> list[tuple[int, int, int]]: ...

def generate_scroll(start_scroll_y, target_scroll_y, *,
                    mode="target", viewport_height=800, seed=None,
                    model_path=None, providers=None) -> list[dict]: ...

• 内部以 (model_path or "__bundled__", tuple(providers or ())) 为 key,functools.lru_cache(maxsize=4) 缓存模型实例
• 多次调用同参数 ⇒ 复用同一 InferenceSession,不重复加载
• maxsize=4 的选择: 典型场景只用 bundled 权重 + 默认 provider(size=1 够用);保留余量给"自定义权重 + CPU/GPU 双 provider"等极少数场景。再大没意义,大模型对象常驻内存反而浪费

与旧 API 的关键差异

旧	新	原因
model_dir: str	model_path: str | Path	接受 Path,语义更明确
config: GenerateConfig 参数	移除	内部物;可调项直接 kwarg
无显式 click 开关	click: bool = True	默认行为不变,但允许纯轨迹
隐式 viewport_norm	viewport_height: int = 800	暴露给调用者
每次调用可能重载模型	类持久 / 函数走 lru_cache	SDK 不能"每次几百 ms 起步"

线程安全

ORT InferenceSession.run 本身线程安全。MouseModel/ScrollModel 共享一个 session,可在多线程并发 .generate()。

seed=None 时使用 np.random.default_rng() 局部实例,不污染全局 numpy 状态。

错误层级

class AiMouseError(Exception): ...
class ModelLoadError(AiMouseError): ...   # 权重缺失 / shape 不符 / providers 不可用
class GenerationError(AiMouseError): ...  # 推理时数值异常 (NaN 等)

§2 ONNX 导出方案

工具脚本 tools/export_onnx.py

一次性脚本(训练完成后跑),产出可烧进 wheel 的 .onnx。库代码不依赖 torch;只有这个脚本依赖。

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

行为:

• 读 train_config.json ⇒ 用一致超参实例化 TrajectoryFlowModel
• torch.load(flow_model.pt) ⇒ model.eval() ⇒ 导出为 flow_model.onnx
• 同样导出 scroll decoder 部分 ⇒ scroll_decoder.onnx(encoder 仅训练用)
• 复制 click_dist.json、duration_dist.json、train_config.json、scroll_config.json 到 assets/

轨迹模型导出

torch.onnx.export(
    model,
    args=(
        torch.zeros(1, seq_len, 3),   # x_t
        torch.zeros(1),               # t
        torch.zeros(1, 3),            # cond
    ),
    f="flow_model.onnx",
    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,
)

关键决策:

• seq_len=64 静态: 位置嵌入是固定 shape 可学习参数;n_points API 若 ≠ 64,在后处理阶段线性插值重采样,不进 ONNX
• Batch 动态: 支持下游批量生成
• Opset 17: 支持 SDPA。如有问题脚本回落到 14
• 常量折叠: 位置嵌入、time/cond MLP 权重可折掉

Euler ODE 循环留在 Python

10 步 Euler 在 Python/NumPy 中循环,每步一次 session.run():

x = rng.standard_normal((1, SEQ_LEN, 3), dtype=np.float32)
dt = 1.0 / N_STEPS
for i in range(N_STEPS):
    t = np.full((1,), (i + 0.5) * dt, dtype=np.float32)
    v = self.session.run(["v"], {"x_t": x, "t": t, "cond": cond_np})[0]
    x = x + dt * v

权衡:10 次 ORT 调用开销 ≈ 1-3 ms;不引入 ONNX Loop op 复杂度;单文件权重;调试容易。未来如需性能可升级单图 Loop 版,本次不做。

滚轮模型导出(只导 decoder)

class ScrollDecoder(nn.Module):
    """Wrap ScrollCVAE.decode for ONNX export."""
    def forward(self, z, cond): ...

torch.onnx.export(
    ScrollDecoder(scroll_cvae),
    args=(torch.zeros(1, LATENT_DIM), torch.zeros(1, 7)),
    f="scroll_decoder.onnx",
    input_names=["z", "cond"],
    output_names=["seq"],
    dynamic_axes={"z": {0: "b"}, "cond": {0: "b"}, "seq": {0: "b"}},
    opset_version=17,
)

z 在 Python 端 rng.standard_normal((1, LATENT_DIM)) 采样;ONNX 内不放采样算子。

体积估算

文件	估算大小
flow_model.onnx	~2-3 MB (4 层 Transformer, d_model=128, ~600K 参数 FP32)
scroll_decoder.onnx	<300 KB
JSON 元数据 (3-4 个)	<20 KB
wheel 总增量	~3 MB

未来可选 --quantize int8 压到 ~1 MB,本次不做(YAGNI)。

兼容性自检(导出脚本内置)

末尾跑烟雾测试:

1. ORT 加载刚导出的 .onnx
2. 与 PyTorch 模型同输入做前向
3. np.allclose(torch_out, ort_out, atol=1e-4) 必须通过
4. 失败 ⇒ 不写入 assets/,保留 .pt

§3 后处理 / 坐标变换的纯 NumPy 重写

coord.py、_gaussian_smooth、_sample_duration、utils.resample_arc、所有空间/时间后处理已经是 numpy。需要替换的只有四处:

现在用	替换为
torch.randn(...)	rng.standard_normal((1, T, 3), dtype=np.float32)
model(x_t, t, cond)	session.run(["v"], {...})[0]
torch.load(...) / .eval()	onnxruntime.InferenceSession(asset_path, providers=...)
scipy.stats.truncnorm.rvs(...)	纯 numpy 拒绝采样(见下)

库内部文件分布

src/ai_mouse/
├── __init__.py          # 公开 API 拼装
├── mouse.py             # MouseModel: session + Euler 循环 + 后处理 + click
├── scroll.py            # ScrollModel: decoder run + 量化 + 时间戳
├── _coord.py            # encode_trajectory, decode_trajectory
├── _postprocess.py      # 平滑、单调性、重采样、截断正态、duration 采样 ...
├── _assets.py           # importlib.resources 加载 onnx + json
├── errors.py            # AiMouseError, ModelLoadError, GenerationError
├── py.typed             # 空文件
└── assets/
    ├── flow_model.onnx
    ├── scroll_decoder.onnx
    ├── click_dist.json
    ├── duration_dist.json
    ├── train_config.json
    └── scroll_config.json

下划线开头的模块是私有实现,不在 __all__,版本间可自由破坏。

tools/ 与 src/ai_mouse/ 的依赖方向: tools/ 是仓库内部的开发期代码,允许 from ai_mouse._coord import ... 反向引用库私有模块(它们在同一 repo 共同演化,只是 tools/ 不进 wheel)。但 src/ai_mouse/ 不允许 import tools.* —— wheel 必须自包含。CI 的 library job 会通过"无 torch 环境"间接守住这条规则。

_truncnorm_sample(替代 scipy)

def _truncnorm_sample(
    mu: float, sigma: float, low: float, high: float,
    rng: np.random.Generator, max_tries: int = 32,
) -> float:
    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))

默认参数 mu=80, sigma=30, low=20, high=300:接受率 ≈ 97%,无可测开销。

_postprocess.py 公开函数清单

def snap_endpoints(forward, lateral, seq_len, n_snap=6) -> tuple[np.ndarray, np.ndarray]: ...
def smooth_start(forward, lateral, n=4) -> tuple[np.ndarray, np.ndarray]: ...
def enforce_forward_monotonic(forward) -> np.ndarray: ...
def gaussian_smooth(x, sigma=1.0) -> np.ndarray: ...
def build_timestamps(log_dt, total_duration_ms, dt_clip=(2.0, 150.0)) -> np.ndarray: ...
def sample_duration(duration_dist, dist, rng) -> float: ...
def truncnorm_sample(mu, sigma, low, high, rng) -> float: ...
def resample_arc(pixels, n_points) -> np.ndarray: ...

每个函数纯函数,显式接收 rng,便于单测。

Scroll 后处理改动

scroll/generator.py 已完全 numpy,迁移只换两行:

• torch.randn(1, latent_dim) ⇒ rng.standard_normal((1, latent_dim), dtype=np.float32)
• model.decode(z, cond_t) ⇒ session.run(["seq"], {"z": z, "cond": cond})[0]

其余 softmax 归一化、量化、时间戳构建照搬。

§4 仓库迁移

终态布局

ai_mouse/                         (repo root)
├── src/
│   └── ai_mouse/                 (唯一进 wheel 的目录)
│       ├── __init__.py
│       ├── mouse.py
│       ├── scroll.py
│       ├── _coord.py
│       ├── _postprocess.py
│       ├── _assets.py
│       ├── errors.py
│       ├── py.typed
│       └── assets/
│           ├── flow_model.onnx
│           ├── scroll_decoder.onnx
│           ├── click_dist.json
│           ├── duration_dist.json
│           ├── train_config.json
│           └── scroll_config.json
├── tools/                        (开发期脚本,不进 wheel)
│   ├── __init__.py
│   ├── __main__.py               (CLI dispatch, 从 ai_mouse/__main__.py 移来)
│   ├── train.py
│   ├── serve.py                  (从 ./main.py 移来)
│   ├── export_onnx.py            ★ 新
│   ├── trainer.py
│   ├── models.py                 (TrajectoryFlowModel — torch)
│   ├── collector.py
│   ├── config.py                 (TrainConfig 等训练侧 dataclass)
│   ├── server/
│   ├── eval/
│   ├── data_adapters/
│   └── scroll/
│       ├── trainer.py
│       ├── models.py
│       └── collector.py
├── data/                         (不变)
├── docs/                         (不变)
├── static/                       (不变,tools/serve.py 引用)
├── tests/
│   ├── unit/                     (库测试,只需 numpy + ort)
│   │   ├── conftest.py
│   │   ├── test_mouse.py
│   │   ├── test_scroll.py
│   │   ├── test_coord.py
│   │   ├── test_postprocess.py
│   │   ├── test_assets.py
│   │   ├── test_errors.py
│   │   └── data/
│   │       ├── golden_mouse.npz
│   │       ├── golden_scroll.npz
│   │       └── tiny_flow.onnx
│   └── tools/                    (训练/服务器测试,需 torch+fastapi+...)
│       ├── conftest.py
│       ├── test_trainer.py
│       ├── test_server.py
│       ├── test_eval_metrics.py
│       ├── test_balabit_adapter.py
│       ├── test_export_onnx.py
│       └── test_scroll_trainer.py
├── examples/
│   └── quickstart.py
├── pyproject.toml                (runtime: numpy+ort;[dev]: torch+fastapi+...)
├── README.md                     (SDK 视角重写)
├── CHANGELOG.md                  (新建,以 0.2.0 起锚)
├── CLAUDE.md                     (更新)
└── uv.lock

被删除:

• ai_mouse/generator.py ⇒ 被 src/ai_mouse/mouse.py 取代
• ai_mouse/scroll/generator.py ⇒ 被 src/ai_mouse/scroll.py 取代
• ai_mouse/utils.py(resample_arc 移入 _postprocess.py)
• JointCVAE(legacy)
• _BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / ... 这类源码树相对路径

迁移分阶段

阶段 1 — 物理移动,训练侧继续工作

• git mv 训练 / 服务器 / 采集 / 评估 / 数据适配器进 tools/
• 修订 tools/ 内部 import 为绝对 from tools.X
• main.py ⇒ tools/serve.py
• tests/test_* 按归属分进 tests/unit/ 或 tests/tools/
• 验收: python -m tools train ...、python -m tools serve 可跑;from ai_mouse import generate 仍能用(老 torch 实现)

阶段 2 — src-layout 切换 + pyproject 更新

• git mv ai_mouse src/ai_mouse
• pyproject 切到 hatchling,运行时依赖收紧到 numpy + onnxruntime;[dev] 群组收纳 torch+fastapi+scipy+matplotlib+pytest+...
• 验收: uv build 产出 wheel,内容仅 src/ai_mouse

阶段 3 — ONNX 导出脚本 + 权重塞进 assets/

• 写 tools/export_onnx.py
• 运行导出,产出 .onnx+元数据 commit 进库
• 验收: pytest tests/tools/test_export_onnx.py 验证 ORT vs PyTorch 数值一致

阶段 4 — 库代码改写(纯 numpy + ORT)

• 抽 _postprocess.py、复制 _coord.py、写 _assets.py、mouse.py、scroll.py、errors.py
• 重写 __init__.py 暴露 API + lru_cache
• 删除阶段 1~2 残留的 torch 引用源文件
• 验收: golden 测试通过;uv pip install . 不引入 torch;from ai_mouse import generate 直接出结果

阶段 5 — 文档与清理

• 重写 README(SDK 视角)
• 新建 CHANGELOG (0.2.0)
• 更新 CLAUDE.md
• 删 JointCVAE
• 写 examples/quickstart.py

风险与对策

风险	对策
权重二进制进 git (~3 MB)	直接 commit;不引入 LFS
中间状态 CI 不绿	每个阶段独立 PR,自包含通过
tools/ 内部 import 风格	一律绝对 from tools.X;tools/__init__.py 留空
阶段 4 改 API 的"半破"中间态	阶段 4 是唯一签名变化节点,同 PR 内同步改 README/CLAUDE/examples
data/models_v2/ 路径引用混乱	库代码不读 data/(只读 importlib.resources);仅 tools/export_onnx.py 和 tools/trainer.py 关心 data/
Golden 测试容差选择	起始 atol=2(像素 + ms),实测后可收紧
Windows 路径	全程 pathlib.Path;importlib.resources 跨平台
ORT CPU/GPU 包互斥	README 明示 pip install onnxruntime-gpu 替换默认 onnxruntime

§5 测试策略

两层测试,依赖边界硬隔离

tests/unit/ 只装 numpy + onnxruntime 必须通过 — "纯推理 SDK"的硬保证。 tests/tools/ 装全套 [dev] 才跑训练/服务器/eval 测试。

Golden 测试(锁住语义不漂移)

迁移开始之前用旧 torch 实现跑 8 个 case × 4 个 seed = 32 个固定输出,存 tests/unit/data/golden_mouse.npz。

CASES = [
    ((100, 200), (900, 400)),   # 水平 800px
    ((500, 500), (500, 100)),   # 垂直 400px
    ((200, 600), (800, 200)),   # 720px 对角
    ((100, 100), (130, 110)),   # 极短 31px
    ((50, 50),   (1500, 900)),  # 极长 1700px
    ((400, 300), (500, 300)),   # 水平 100px
    ((300, 300), (700, 700)),   # 45° 对角
    ((600, 400), (200, 100)),   # 反向对角
]
SEEDS = (0, 1, 2, 3)

迁移后断言 np.allclose(new_pts, golden_pts, atol=2)(像素 + ms 容差)。

Scroll 同样做 golden;容差以 deltaY 量化粒度允许 ±1 个 quantum。

ONNX vs PyTorch 数值一致性

tests/tools/test_export_onnx.py:训练 1-epoch toy 模型(seq_len=8, d_model=16)⇒ 导出 ⇒ 同输入跑 PyTorch & ORT ⇒ assert np.allclose(atol=1e-4)。

真模型的一致性由 tools/export_onnx.py 末尾烟雾测试在导出时执行。

_postprocess.py 单元测试

每个纯函数单独测,边界条件覆盖。详细测试矩阵见原讨论;关键覆盖点:

函数	关键断言
snap_endpoints	输出端点严格 = (0,0) 和 (1,0);前 (seq_len - n_snap) 个点不变
enforce_forward_monotonic	输出非递减
gaussian_smooth	端点保持;权重和归一
build_timestamps	sum ≈ total_duration;严格单调;首项 0
sample_duration	1000 次采样中位数落在 bin 对应 [0.5×, 2×]exp(mu_log)
truncnorm_sample	落在 [low, high];均值 ≈ mu;兜底分支独立覆盖
resample_arc	长度变换正确;N=M 恒等

MouseModel / ScrollModel 行为测试

• session 复用(同实例多次 generate 不重建 session)
• 顶层 generate() 单例缓存命中
• seed 复现性(同 seed 两次调用结果一致)
• click=False 返回纯轨迹
• 非法 model_path 抛 ModelLoadError

线程安全测试

ThreadPoolExecutor 并发 32 次 m.generate(...),断言全部完成、无异常。不保证并发下同 seed 严格一致(ORT 内部调度)。

CI 矩阵

jobs:
  library:
    # 只装运行时依赖 + pytest — 拦截库代码偷偷 import torch
    runs: |
      uv venv && uv pip install -e . pytest
      uv run pytest tests/unit -v
  dev:
    runs: |
      uv sync --group dev
      uv run pytest tests/ -v

注: pytest 不能进运行时依赖(下游用户不需要),只能在 CI 安装行里临时加。 pytest-asyncio 仅 tests/tools/test_server.py 需要,留在 [dev] 群组。

OS 矩阵: Linux + Windows;Python 3.12 + 3.13。

不做的事

• ❌ 不测视觉真实性(主观,留给人眼审 eval report)
• ❌ 不跑真 web E2E(httpx ASGI 集成已足够)
• ❌ CI 不重训模型(用已 commit 的 ONNX)

§6 文档与下游集成

README.md 重写(SDK 视角)

读者画像从"训练者"切换到"集成者"。结构:

1. 标题与一句话定位
2. 安装(pip install git+... + GPU 可选说明)
3. 快速上手(鼠标、滚轮、类用法、provider 切换、复现性)
4. API 概览(表格)
5. 线程安全说明
6. 自训练模型 — 跳转 Development 小节
7. Development(训练命令、Web UI、ONNX 导出 — 目标读者是 contributor)

examples/quickstart.py

最小可跑示例:from ai_mouse import generate ⇒ 打印 6 行结果 + 一段"实际使用模式"伪代码(时间戳节流 + 占位的鼠标移动调用)。

CHANGELOG.md(新建,Keep a Changelog 格式)

锚定 0.2.0:

## [0.2.0] - YYYY-MM-DD
### Changed (breaking)
- Inference 不再依赖 PyTorch;运行时仅 numpy + onnxruntime
- 公开 API 新增 MouseModel / ScrollModel 类
- 函数式 generate / generate_scroll 签名: keyword-only;新增 click=, seed=;移除 config=
- 训练 / 服务器 / 采集 / 评估代码移至 tools/,不再打入 wheel
### Added
- Wheel 内置 ONNX 预训练权重 (~3 MB)
- ORT provider 自定义 (GPU/DirectML)
### Removed
- JointCVAE (legacy)
- ai_mouse.config.GenerateConfig 顶层导出

CLAUDE.md 更新清单

• Commands: python -m ai_mouse <cmd> ⇒ python -m tools <cmd>
• 加一节"Library vs tools 边界": 库代码禁止 import torch
• 移除 "Bundled weights live in data/models_v2/" ⇒ 改 src/ai_mouse/assets/
• 测试章节区分 tests/unit / tests/tools

下游契约 — 成功验收标准

1. pip install git+... ⇒ from ai_mouse import generate 直接出结果,无缺失文件、无 GB 级依赖
2. import ai_mouse 时间 < 200 ms
3. 单次 generate() < 50 ms (CPU,Euler 10 步 + 后处理)

不做的事

• ❌ 不写 API reference 网站(docstring 足够)
• ❌ 不写迁移 guide(破坏面有限,CHANGELOG 一段话)
• ❌ 不录视频 / GIF
• ❌ 不引入 docs i18n 流程

开放问题

无。所有决策点已经在与用户的对话中明确。

后续行动

本 spec 经用户批准后,进入 implementation plan 阶段(writing-plans skill),把 §4 的 5 个迁移阶段展开为可执行的 task list。