streamembed 0.1.0__py3-none-any.whl

streamembed/__init__.py

@@ -0,0 +1,7 @@
+"""StreamEmbed: search your own screen history by describing what you saw. Everything runs on your machine."""
+from .api import (Capture, Result, available_index_types, available_models, build_index_from_directory, build_index_from_now, configure, query_index,
+                  visualize_results)
+
+__version__ = "0.1.0"
+__all__ = ["Capture", "Result", "available_index_types", "available_models", "build_index_from_directory", "build_index_from_now", "configure", "query_index",
+           "visualize_results"]

streamembed/api.py

@@ -0,0 +1,205 @@
+"""Index screenshots from a folder or from the screen, query them, look at the results.
+
+    import streamembed as se
+    se.build_index_from_directory("screenshots")
+    capture = se.build_index_from_now(interval=10)          # keeps capturing in the background
+    results = se.query_index("the flight search page with fares to Denver", top_k=10)
+    se.visualize_results(results)
+
+Everything acts on one index folder (default ./streamembed_index); `configure()` changes it and the other settings.
+"""
+from __future__ import annotations
+
+import threading
+import time
+import webbrowser
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterator
+
+from .capture import folder_frames, screen_frames, screen_permission
+from .gate import Gate
+from .index import INDEX_TYPES, Index
+from . import models as _models
+from .models import DEFAULT_MODEL, resolve
+
+JPEG_QUALITY = 85
+
+
+@dataclass
+class Result:
+    id: int             # row in the index
+    score: float        # cosine similarity
+    time: str           # capture time, local, "YYYY-MM-DD HH:MM:SS"
+    ts: float           # the same, unix seconds
+    app: str | None
+    title: str | None   # window title; for folder screenshots, the source file
+    path: str           # the stored JPEG
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class Capture:
+    """Screen capture in the background: `stop()`, `pause()`, `resume()`, `wait()`, and `stats` (captured, encoded, idle, skipped, errors)."""
+
+    def __init__(self, interval: float, duration: float | None, frames: Iterator):
+        self.stats = dict(captured=0, encoded=0, idle=0, skipped=0, errors=0)
+        self._stop, self._pause = threading.Event(), threading.Event()
+        self._thread = threading.Thread(target=_ingest, args=(frames, self.stats, self._stop, self._pause, interval), daemon=True)
+        self._thread.start()
+        if duration:
+            threading.Timer(duration, self.stop).start()
+
+    @property
+    def running(self) -> bool:
+        return self._thread.is_alive()
+
+    def pause(self) -> None:
+        self._pause.set()
+
+    def resume(self) -> None:
+        self._pause.clear()
+
+    def stop(self) -> None:
+        self._stop.set()
+        self._thread.join(timeout=30)
+
+    def wait(self) -> None:
+        """Block until the capture stops (after `duration`, or Ctrl-C)."""
+        try:
+            while self._thread.is_alive():
+                self._thread.join(1.0)
+        except KeyboardInterrupt:
+            self.stop()
+
+
+# ---- settings and the lazily loaded pieces (the model takes ~20 s and ~8 GB, so it loads once, on first use)
+_cfg = dict(index_dir=Path("streamembed_index"), index_type="flat", model=DEFAULT_MODEL, models_repo=None, device="auto", width=1280)
+_state: dict = dict(encoder=None, index=None, capture=None, query=None)
+
+
+def available_index_types() -> list[str]:
+    return list(INDEX_TYPES)
+
+
+def available_models() -> list[str]:
+    """Model names from the configured model repo, then the built-in base models."""
+    return _models.available_models(_cfg["models_repo"])
+
+
+def configure(index_dir: str | Path = "streamembed_index", index_type: str = "flat", model: str = DEFAULT_MODEL,
+              models_repo: str | Path | None = None, device: str = "auto", width: int = 1280, encoder=None) -> None:
+    """index_dir: where frames, vectors and images go. index_type: "flat" (exact) or "binary" (1-bit codes + re-rank).
+    model: a name from `available_models()`. models_repo: the Hugging Face repo or local folder with models.json (default
+    $STREAMEMBED_MODELS, else the StreamEmbed repo). encoder: a stand-in, for tests."""
+    if index_type not in INDEX_TYPES:
+        raise ValueError(f"Invalid index type: {index_type} not in {available_index_types()}")
+    if _state["capture"] is not None:
+        _state["capture"].stop()
+    _cfg.update(index_dir=Path(index_dir), index_type=index_type, model=model, models_repo=str(models_repo) if models_repo else None,
+                device=device, width=width)
+    _state.update(encoder=encoder, index=None, capture=None, query=None)
+
+
+def _encoder():
+    if _state["encoder"] is None:
+        from .encoder import Encoder
+        spec = resolve(_cfg["model"], _cfg["models_repo"])
+        _state["encoder"] = Encoder(spec.weights, _cfg["device"], spec.base, spec.revision, (spec.scale, spec.scale), spec.encoder_name)
+    return _state["encoder"]
+
+
+def _index() -> Index:
+    if _state["index"] is None:
+        _state["index"] = Index(_cfg["index_dir"], _cfg["index_type"], _encoder().name)
+    return _state["index"]
+
+
+def _ingest(frames: Iterator, stats: dict, stop: threading.Event | None = None, pause: threading.Event | None = None,
+            interval: float | None = None) -> dict:
+    """frames -> change gate -> JPEG -> stream encoder -> index. One call is one stream."""
+    enc, idx, gate = _encoder(), _index(), Gate()
+    if hasattr(enc, "reset_stream"):
+        enc.reset_stream()          # nothing is reused from the previous source's frames
+    it = iter(frames)
+    while stop is None or not stop.is_set():
+        tick = time.monotonic()
+        if pause is not None and pause.is_set():      # checked before taking a frame: nothing is captured while paused
+            stop.wait(0.5)
+            continue
+        try:
+            item = next(it)
+        except StopIteration:
+            break
+        try:
+            if item is None:
+                stats["skipped"] += 1
+            else:
+                img, ts, app, title = item
+                stats["captured"] += 1
+                d = gate(img, ts)
+                if d.idle:
+                    stats["idle"] += 1
+                    idx.still_on_screen(ts)
+                if d.encode:
+                    rel = Path("images") / time.strftime("%Y-%m-%d", time.localtime(ts)) / f"{int(ts * 1000)}_{len(idx)}.jpg"
+                    (idx.root / rel).parent.mkdir(parents=True, exist_ok=True)
+                    img.save(idx.root / rel, "JPEG", quality=JPEG_QUALITY)
+                    vec, _ = enc.encode_frame(idx.root / rel)
+                    idx.add(vec, dict(ts=ts, path=str(rel), app=app, title=title))
+                    stats["encoded"] += 1
+        except Exception:
+            stats["errors"] += 1
+        if interval and stop is not None:
+            stop.wait(max(0.0, interval - (time.monotonic() - tick)))
+    return stats
+
+
+# ---- the API
+def build_index_from_directory(directory: str | Path, order: str = "mtime") -> dict:
+    """Index every screenshot in `directory` (searched recursively) as one stream. order="mtime" uses each file's time as its capture
+    time; order="name" sorts by path. Unchanged screens are skipped. Returns the counts."""
+    if _state["capture"] is not None and _state["capture"].running:
+        raise RuntimeError("Stop the running capture first: both feed one stream encoder.")
+    return _ingest(folder_frames(Path(directory), _cfg["width"], order), dict(captured=0, encoded=0, idle=0, skipped=0, errors=0))
+
+
+def build_index_from_now(interval: float = 5.0, duration: float | None = None, block: bool = False, frames: Iterator | None = None) -> Capture:
+    """Capture the primary display every `interval` seconds from now on, in the background (or until it stops, with block=True).
+    Skipped while the screen is locked or a password manager or private browsing window is in front. `frames` replaces the screen (tests)."""
+    if _state["capture"] is not None and _state["capture"].running:
+        raise RuntimeError("A capture is already running.")
+    if frames is None:
+        if not screen_permission():
+            raise PermissionError("Screen Recording permission is missing. Allow it for this app in System Settings > Privacy & Security "
+                                  "> Screen Recording, then run again.")
+        frames = screen_frames(_cfg["width"])
+    _encoder(), _index()                   # load before the thread starts, so errors surface here
+    _state["capture"] = Capture(interval, duration, frames)
+    if block:
+        _state["capture"].wait()
+    return _state["capture"]
+
+
+def query_index(query: str, top_k: int = 10) -> list[Result]:
+    """Screens that best match a description. A match within 30 s of a better one is dropped."""
+    idx = _index()
+    hits = idx.search(_encoder().encode_query(query), top_k)
+    _state["query"] = query
+    out = []
+    for row, score in hits:
+        f = idx.frames[row]
+        out.append(Result(row, round(score, 4), time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(f["ts"])), f["ts"], f["app"],
+                          f["title"], str((idx.root / f["path"]).resolve())))
+    return out
+
+
+def visualize_results(results: list[Result], path: str | Path | None = None, open_browser: bool = True, context: int = 15) -> Path:
+    """Write an HTML page of the results and open it. Each result opens full size with the `context` frames before and after it."""
+    from .visualize import write_page
+    idx = _index()
+    out = write_page(idx, results, Path(path) if path else idx.root / "results.html", context, _state["query"])
+    if open_browser:
+        webbrowser.open(out.resolve().as_uri())
+    return out

streamembed/capture.py

@@ -0,0 +1,88 @@
+"""Frames to index: the primary display every few seconds, or a folder of screenshots. Each frame is (image, ts, app, title)."""
+from __future__ import annotations
+
+import sys
+import time
+from pathlib import Path
+from typing import Iterator
+
+from PIL import Image
+
+BLOCK = 32
+EXTS = {".jpg", ".jpeg", ".png", ".webp"}
+# Not captured when the frontmost app or window title contains one of these (case-insensitive). Other windows on screen are not checked.
+EXCLUDE_APPS = ("1password", "bitwarden", "dashlane", "lastpass", "keychain access", "keepass")
+EXCLUDE_TITLES = ("private browsing", "incognito", "inprivate")
+MACOS = sys.platform == "darwin"
+
+
+def fit_width(img: Image.Image, width: int) -> Image.Image:
+    """Resize to `width` keeping the aspect ratio, both sides multiples of 32 (one block is 32x32 px)."""
+    w = max(BLOCK, width // BLOCK * BLOCK)
+    h = max(BLOCK, round(img.height * w / img.width / BLOCK) * BLOCK)
+    img = img.convert("RGB")
+    return img if img.size == (w, h) else img.resize((w, h), Image.BICUBIC)
+
+
+def folder_frames(folder: Path, width: int, order: str = "mtime", spacing_s: float = 5.0) -> Iterator[tuple]:
+    """order="mtime": by file time, and each file's time is its capture time. order="name": by path, stamped `spacing_s` apart, ending now."""
+    if order not in ("mtime", "name"):
+        raise ValueError(f"order must be 'mtime' or 'name', not {order!r}")
+    files = [p for p in Path(folder).rglob("*") if p.suffix.lower() in EXTS and p.is_file()]
+    if not files:
+        raise FileNotFoundError(f"No screenshots (.jpg, .jpeg, .png, .webp) found in {folder}")
+    files.sort(key=(lambda p: (p.stat().st_mtime, str(p))) if order == "mtime" else str)
+    start = time.time() - len(files) * spacing_s
+    for i, p in enumerate(files):
+        with Image.open(p) as im:
+            yield fit_width(im, width), p.stat().st_mtime if order == "mtime" else start + i * spacing_s, None, str(p)
+
+
+# ---- the screen. macOS reports permission, lock state and the front window; elsewhere these checks pass.
+def screen_permission() -> bool:
+    """Without Screen Recording permission macOS returns only the wallpaper."""
+    if not MACOS:
+        return True
+    import Quartz
+    if Quartz.CGPreflightScreenCaptureAccess():
+        return True
+    Quartz.CGRequestScreenCaptureAccess()      # shows the system prompt
+    return False
+
+
+def _locked() -> bool:
+    if not MACOS:
+        return False
+    import Quartz
+    session = Quartz.CGSessionCopyCurrentDictionary()
+    return bool(session and session.get("CGSSessionScreenIsLocked", 0))
+
+
+def _front_window() -> tuple[str | None, str | None]:
+    """(app, window title) of the frontmost normal window."""
+    if not MACOS:
+        return None, None
+    import Quartz
+    options = Quartz.kCGWindowListOptionOnScreenOnly | Quartz.kCGWindowListExcludeDesktopElements
+    for win in Quartz.CGWindowListCopyWindowInfo(options, Quartz.kCGNullWindowID) or []:
+        if win.get("kCGWindowLayer", 1) == 0:
+            return win.get("kCGWindowOwnerName"), win.get("kCGWindowName") or None
+    return None, None
+
+
+def private(app: str | None, title: str | None) -> bool:
+    a, t = (app or "").casefold(), (title or "").casefold()
+    return any(x in a for x in EXCLUDE_APPS) or any(x in t for x in EXCLUDE_TITLES)
+
+
+def screen_frames(width: int) -> Iterator[tuple | None]:
+    """One item per call: a frame, or None while the screen is locked or a private window is in front."""
+    import mss
+    with mss.mss() as sct:
+        while True:
+            app, title = _front_window()
+            if _locked() or private(app, title):
+                yield None
+                continue
+            shot = sct.grab(sct.monitors[1])
+            yield fit_width(Image.frombytes("RGB", shot.size, shot.rgb), width), time.time(), app, title

streamembed/encoder.py

@@ -0,0 +1,108 @@
+"""One model, two sets of LLM attention weights.
+
+Frames are encoded with the consistency adapter and questions with the screenshot LoRA. Both are the same base model plus a low-rank update
+of the same attention layers, so instead of loading the model twice the LLM attention weights are held twice and swapped by reference for
+the length of a question. The ViT is not stored twice because a text question never runs it. Weights are computed exactly as `merge_lora`
+does (W + B @ A in fp32, cast back), so results equal two separately loaded models.
+"""
+from __future__ import annotations
+
+import threading
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .engine.config import MODEL_ID, PRESETS, QUERY_INSTRUCTION
+from .engine.engine import Engine, preprocess
+from .engine.model import Model, _lora_pairs, pick_device
+from .engine.qwen_embedder import Qwen3VLEmbedder
+
+METHOD = "S+G+T+L"      # the paper's method: spatial pruning with Grid keep + temporal reuse (ViT and LLM)
+
+
+def _updated(w: torch.Tensor, a: torch.Tensor, b: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+    return (w.data.float() + scale * (b.float() @ a.float())).to(w.dtype)
+
+
+class Encoder:
+    def __init__(self, weights: tuple[Path, Path] | None, device: str = "auto", model_id: str = MODEL_ID, revision: str | None = None,
+                 scales: tuple[float, float] = (1.0, 1.0), name: str | None = None):
+        """weights: (frame adapter, question LoRA) PEFT state dicts, or None for the base model. revision: the base model commit.
+        scales: LoRA alpha / rank of each."""
+        device = pick_device(device)
+        dtype = torch.float32 if device == "cpu" else torch.float16
+        try:
+            emb = Qwen3VLEmbedder(model_id, dtype=dtype, local_files_only=True, revision=revision)     # once cached, never touch the network
+        except OSError:
+            emb = Qwen3VLEmbedder(model_id, dtype=dtype, revision=revision)                            # first run: download
+        query_weights = self._merge(emb.model, *weights, *scales) if weights is not None else []
+        emb.model.to(device).eval()
+        # Read the frame weights after .to(device), which may replace the Parameter objects
+        self._swap = [(mod, mod.weight, torch.nn.Parameter(q.to(device), requires_grad=False)) for mod, q in query_weights]
+        self.m = Model(embedder=emb, device=torch.device(device), dtype=dtype)
+        self.cfg = PRESETS[METHOD]
+        self.engine = Engine(self.m, self.cfg)
+        self.lock = threading.Lock()    # one forward at a time: the swap must not overlap a frame
+        self._first = True
+        self.name = name or f"{Path(model_id).name}:{weights[0].name if weights else 'zero-shot'}:{METHOD}"
+
+    @staticmethod
+    def _merge(model: torch.nn.Module, adapter: Path, query: Path, frame_scale: float = 1.0,
+               query_scale: float = 1.0) -> list[tuple[torch.nn.Module, torch.Tensor]]:
+        """Merge the adapter into the model; return (module, query weight) for every LLM layer either LoRA touches."""
+        frame_pairs = {name: (a, b) for name, a, b in _lora_pairs(str(adapter))}
+        query_pairs = {name: (a, b) for name, a, b in _lora_pairs(str(query)) if "language_model" in name}
+        out = []
+        for name in sorted(frame_pairs.keys() | query_pairs.keys()):
+            mod = model.get_submodule(name)
+            w = mod.weight
+            if "language_model" in name:
+                out.append((mod, _updated(w, *query_pairs[name], query_scale) if name in query_pairs else w.data.clone()))
+            if name in frame_pairs:
+                w.data = _updated(w, *frame_pairs[name], frame_scale)
+        return out
+
+    @contextmanager
+    def _query_weights(self):
+        for mod, _, q in self._swap:
+            mod.weight = q
+        try:
+            yield
+        finally:
+            for mod, f, _ in self._swap:
+                mod.weight = f
+
+    def reset_stream(self) -> None:
+        """The next frame starts a new stream: nothing is reused from frames of another source."""
+        with self.lock:
+            self._first = True
+
+    def encode_frame(self, image_path: str | Path) -> tuple[np.ndarray, float]:
+        """Next frame of the stream -> (2048-d unit vector, encode milliseconds)."""
+        t = time.perf_counter()
+        frame = preprocess(self.m, str(image_path), self.cfg)
+        with self.lock:
+            emb, _ = self.engine.step(frame, first=self._first)
+            self._first = False
+            v = emb.float().cpu().numpy()
+            self._trim()
+        return v, (time.perf_counter() - t) * 1000
+
+    def _trim(self) -> None:
+        """Metal keeps freed buffers unless asked to release them. On an M2 Pro over 245 frames, releasing after every frame cut peak Metal
+        memory from 10.3 GB (every 50 frames) to 7.1 GB, with no slowdown (mean 675 vs 695 ms)."""
+        if self.m.device.type == "mps":
+            torch.mps.empty_cache()
+
+    def encode_query(self, text: str) -> np.ndarray:
+        with self.lock, self._query_weights():
+            self.m.core.rope_deltas = None      # cached by the stock forward; must not leak between prompts
+            e = self.m.embedder.embed([self.m.embedder.format_model_input(text=text, instruction=QUERY_INSTRUCTION)], normalize=False)
+            self.m.core.rope_deltas = None
+            v = F.normalize(e.float(), p=2, dim=-1)[0].cpu().numpy()
+            self._trim()
+        return v

streamembed/engine/__init__.py

@@ -0,0 +1 @@
+"""Stream encoder from the paper: spatial pruning (with Grid keep) and temporal reuse. Code names follow the paper code (Token Pruning, KV Reuse)."""

streamembed/engine/config.py

@@ -0,0 +1,82 @@
+"""Method settings: which efficiency components are on, their thresholds, and named presets.
+
+Pixel values from the Qwen3-VL processor are normalised to [-1, 1], so one 8-bit grey level is 1/127.5. Thresholds are given in grey
+levels on the command line (tau 8, delta 8) and converted with LEVEL.
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, replace
+
+MODEL_ID = os.environ.get("STREAMEMBED_MODEL", "Qwen/Qwen3-VL-Embedding-2B")
+LEVEL = 1 / 127.5                       # one 8-bit grey level in processor units
+QUERY_INSTRUCTION = "Retrieve the computer screen that matches the user's description."
+MAX_OCR_CHARS = 12000                   # OCR text is cut to this many characters before text embedding
+
+
+@dataclass(frozen=True)
+class MethodConfig:
+    """Efficiency components of the screenshot route.
+
+    spatial             Token Pruning rule: "none" | "lt" (drop a block equal within tau to its left OR upper block) | "lt-grid"
+                        (paper: lt + Grid Keep, one block per flat region per grid_tile x grid_tile block tile kept) | "pixelprune" (Pred-2D)
+    tau_levels          Token Pruning threshold in grey levels (8 = 8/255)
+    grid_tile           Grid Keep tile side in blocks (paper: 3, i.e. 3x3-block tiles)
+    temporal_vit        KV Reuse (ViT): recompute only changed blocks in the ViT (cached K/V for the rest); implies exact LLM prefix reuse
+    delta_levels        KV Reuse change threshold in grey levels (block changed if max |pixel - pixel at its last recompute| > delta)
+    refresh_rho         KV Reuse: full refresh when more than this fraction of the kept blocks changed ...
+    refresh_every       ... or every K-th frame since the last refresh
+    approx_llm          KV Reuse (LLM): recompute only changed visual tokens + text after the first change (needs temporal_vit)
+    linear_patch_embed  linear patch embedding ("+pe"): ViT patch embedding as F.linear instead of Conv3d (identical math; MPS Conv3d
+                        is very slow)
+    """
+    spatial: str = "none"
+    tau_levels: float = 8.0
+    grid_tile: int = 3
+    temporal_vit: bool = False
+    delta_levels: float = 8.0
+    refresh_rho: float = 0.9
+    refresh_every: int = 10
+    approx_llm: bool = False
+    linear_patch_embed: bool = False
+
+    def __post_init__(self):
+        assert self.spatial in ("none", "lt", "lt-grid", "pixelprune"), self.spatial
+        assert self.temporal_vit or not self.approx_llm, "KV Reuse (LLM) reuses per-block state that only KV Reuse (ViT) keeps"
+        assert self.grid_tile >= 1, self.grid_tile
+
+    @property
+    def tau(self) -> float:
+        return self.tau_levels * LEVEL
+
+    @property
+    def delta(self) -> float:
+        return self.delta_levels * LEVEL
+
+    @property
+    def name(self) -> str:
+        """Short tag, e.g. LT8t3TL+pe (Token Pruning rule + tau, t3 = Grid Keep with 3x3 tiles, T = KV Reuse (ViT), L = KV Reuse
+        (LLM), +pe = linear patch embedding)."""
+        s = {"none": "FULL", "lt": f"LT{self.tau_levels:g}", "lt-grid": f"LT{self.tau_levels:g}t{self.grid_tile}",
+             "pixelprune": f"PP{self.tau_levels:g}"}[self.spatial]
+        if self.spatial == "none" and self.temporal_vit:
+            s = ""
+        s += ("T" if self.temporal_vit else "") + ("L" if self.approx_llm else "")
+        if self.temporal_vit and (self.refresh_rho, self.refresh_every, self.delta_levels) != (0.9, 10, 8.0):
+            s += f"_r{self.refresh_rho:g}k{self.refresh_every}d{self.delta_levels:g}"
+        return s + ("+pe" if self.linear_patch_embed else "")
+
+    def with_(self, **kw) -> "MethodConfig":
+        return replace(self, **kw)
+
+
+# Named presets (paper rows). Explicit CLI flags override preset values.
+PRESETS: dict[str, MethodConfig] = {
+    "full": MethodConfig(),                                                             # Full screenshot, stock Conv3d patch embedding
+    "full+pe": MethodConfig(linear_patch_embed=True),                                   # Full screenshot, linear patch embedding
+    "S": MethodConfig(spatial="lt", linear_patch_embed=True),                           # Token Pruning
+    "S+T": MethodConfig(spatial="lt", temporal_vit=True, linear_patch_embed=True),      # + KV Reuse (ViT) (+ exact LLM prefix reuse)
+    "S+T+L": MethodConfig(spatial="lt", temporal_vit=True, approx_llm=True, linear_patch_embed=True),   # + KV Reuse (LLM)
+    "S+G": MethodConfig(spatial="lt-grid", linear_patch_embed=True),                    # Token Pruning with Grid Keep
+    "S+G+T+L": MethodConfig(spatial="lt-grid", temporal_vit=True, approx_llm=True, linear_patch_embed=True),   # paper method
+}

streamembed/engine/engine.py

@@ -0,0 +1,90 @@
+"""Engine: one screenshot in, one embedding out, carrying per-stream state between frames.
+
+Engine.step:  Token Pruning keep set + KV Reuse (ViT) plan (change detection, refresh rule)
+              ->  LLM plan (prefix P, skip, KV Reuse (LLM) rows R)
+              ->  ViT (full or partial)  ->  sync  ->  LLM (exact prefix reuse or KV Reuse (LLM))  ->  stats / MAC row  ->  commit state
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+
+from . import llm
+from .config import MethodConfig
+from .macs import llm_macs_rows, llm_macs_suffix, vit_macs
+from .model import Grid, Model, grid
+from .spatial import keep_mask
+from .temporal import FramePlan, StreamState, commit, plan_frame
+from .vit import run_vit
+
+
+@dataclass
+class FrameInputs:
+    inputs: dict                        # processor output (CPU tensors)
+    keep_mask: Optional[np.ndarray]     # Token Pruning keep mask (None = all blocks)
+
+
+def preprocess(m: Model, image_path: str, cfg: MethodConfig) -> FrameInputs:
+    """Image file -> processor tensors (+ Token Pruning keep mask). CPU only, so it can run in DataLoader workers."""
+    emb = m.embedder
+    inputs = emb.preprocess([emb.format_model_input(image=image_path)])
+    return FrameInputs(inputs, keep_mask(cfg.spatial, inputs["pixel_values"], inputs["image_grid_thw"], cfg.tau, cfg.grid_tile))
+
+
+class Engine:
+    """One engine per stream and config (own caches)."""
+
+    def __init__(self, m: Model, cfg: MethodConfig):
+        self.m, self.cfg = m, cfg
+        self.state: Optional[StreamState] = None
+
+    @torch.inference_mode()
+    def step(self, frame: FrameInputs, first: bool, timings: Optional[dict] = None) -> tuple[torch.Tensor, dict]:
+        """Embed one frame. first=True starts a new stream (drops all reuse state). Returns (embedding (2048,), stats row)."""
+        m, cfg = self.m, self.cfg
+        inputs = {k: v.to(m.device) for k, v in frame.inputs.items()}
+        g = grid(m, inputs)
+        pixels = inputs["pixel_values"].float()
+        block_pixels = pixels.reshape(g.n_blocks, -1)
+
+        state = None
+        if cfg.temporal_vit:
+            if self.state is None or self.state.key != g.key:   # new resolution / prompt: fresh caches
+                self.state = StreamState(m, g, approx_llm=cfg.approx_llm)
+                first = True
+            state = self.state
+            if first:
+                state.reset(m, g)
+
+        plan = plan_frame(cfg, m, g, state, block_pixels, first, frame.keep_mask)
+        llm.plan_llm(g, plan, cfg.approx_llm)
+        visual_tokens, deepstack = run_vit(m, g, plan, pixels, block_pixels, cfg.linear_patch_embed)
+        m.sync()
+        if timings is not None:
+            timings["t_vit_done"] = time.perf_counter()
+        run_llm = llm.run_approx if cfg.approx_llm else llm.run_exact
+        emb = run_llm(m, g, plan, visual_tokens, deepstack)
+        row = stats_row(m, g, plan, cfg.approx_llm)
+        commit(plan, emb)
+        return emb, row
+
+
+def stats_row(m: Model, g: Grid, plan: FramePlan, approx_llm: bool) -> dict:
+    """Per-frame work counters and MACs."""
+    nk, nc = len(plan.kept), len(plan.recompute)
+    L = g.llm_len(nk)
+    P = L if plan.skip_llm else plan.prefix
+    if plan.skip_llm:
+        n_rows, lm = 0, 0
+    elif approx_llm:
+        n_rows, lm = len(plan.llm_rows), llm_macs_rows(m, plan.llm_rows)
+    else:
+        n_rows, lm = L - P, llm_macs_suffix(m, P, L)
+    return dict(n_kept=nk, n_comp_blocks=nc, full=bool(plan.refresh), changed_frac=plan.changed_frac,
+                vit_q_patches=4 * nc, vit_kv_patches=4 * nk, vit_macs=vit_macs(m, 4 * nc, 4 * nk, nc) if nc else 0,
+                llm_comp_tokens=n_rows, llm_len=L, llm_macs=lm, since_full=plan.since_refresh, llm_skip=bool(plan.skip_llm),
+                llm_P=int(P))