From 4e69ecc963e1b2ffb7111bebeea8987b71e72aee Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Tue, 12 May 2026 01:14:37 +0800 Subject: [PATCH] feat(lib): add ScrollModel (numpy + ONNX Runtime); rename legacy scroll module --- src/ai_mouse/__init__.py | 2 +- .../generator.py => _scroll_legacy.py} | 0 src/ai_mouse/scroll.py | 139 ++++++++++++++++++ src/ai_mouse/scroll/__init__.py | 4 - tests/unit/test_scroll.py | 37 +++++ tests/unit/test_scroll_generator.py | 2 +- 6 files changed, 178 insertions(+), 6 deletions(-) rename src/ai_mouse/{scroll/generator.py => _scroll_legacy.py} (100%) create mode 100644 src/ai_mouse/scroll.py delete mode 100644 src/ai_mouse/scroll/__init__.py create mode 100644 tests/unit/test_scroll.py diff --git a/src/ai_mouse/__init__.py b/src/ai_mouse/__init__.py index f0e1957..027a92a 100644 --- a/src/ai_mouse/__init__.py +++ b/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"] diff --git a/src/ai_mouse/scroll/generator.py b/src/ai_mouse/_scroll_legacy.py similarity index 100% rename from src/ai_mouse/scroll/generator.py rename to src/ai_mouse/_scroll_legacy.py diff --git a/src/ai_mouse/scroll.py b/src/ai_mouse/scroll.py new file mode 100644 index 0000000..52d8812 --- /dev/null +++ b/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() diff --git a/src/ai_mouse/scroll/__init__.py b/src/ai_mouse/scroll/__init__.py deleted file mode 100644 index 8dd7bfa..0000000 --- a/src/ai_mouse/scroll/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Scroll wheel event generation (inference only).""" -from ai_mouse.scroll.generator import generate_scroll - -__all__ = ["generate_scroll"] diff --git a/tests/unit/test_scroll.py b/tests/unit/test_scroll.py new file mode 100644 index 0000000..93eb36c --- /dev/null +++ b/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") diff --git a/tests/unit/test_scroll_generator.py b/tests/unit/test_scroll_generator.py index e4cc098..b143433 100644 --- a/tests/unit/test_scroll_generator.py +++ b/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