feat(lib): add ScrollModel (numpy + ONNX Runtime); rename legacy scroll module

src/ai_mouse/__init__.py

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

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()

src/ai_mouse/scroll/__init__.py

@@ -1,4 +0,0 @@
-"""Scroll wheel event generation (inference only)."""
-from ai_mouse.scroll.generator import generate_scroll
-
-__all__ = ["generate_scroll"]

tests/unit/test_scroll.py

@@ -0,0 +1,37 @@
+"""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")

tests/unit/test_scroll_generator.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import torch
 import pytest
 
-from ai_mouse.scroll.generator import generate_scroll
+from ai_mouse._scroll_legacy import generate_scroll
 from tools.scroll.models import ScrollCVAE