streamembed 0.1.0__py3-none-any.whl

streamembed/engine/spatial.py

@@ -0,0 +1,126 @@
+"""Token Pruning, for spatial redundancy: drop 32x32-px blocks (one LLM token each) before the ViT, within a frame.
+
+Much of a screen is blank or flat (backgrounds, margins, whitespace, table fills). A block whose pixels are predictable from its left or
+upper neighbour carries little information, so it enters neither the ViT nor the LLM. Rules over the processor's pixel values (tau 8
+grey levels = 8/255):
+
+  lt          (Token Pruning, "LT8") drop a block iff every pixel is within tau of its LEFT or its UPPER neighbour block, comparing
+              original pixels (order-free version of PixelPrune). Every block is decided independently -> a few vectorised array ops.
+  lt-grid     (Token Pruning with Grid Keep, "LT8t3", paper default) lt alone shrinks a large flat region (blank page, empty panel) to a
+              single token, so the model loses its size and layout. Grid Keep keeps one block per tile x tile block tile of each flat
+              region (paper: 3x3; MethodConfig.grid_tile), at its own position. No training; kept blocks are a superset of lt's.
+  pixelprune  PixelPrune Pred-2D: raster-order scan with a lossless-coding predictor from the reconstructed left / up / up-left
+              neighbours (kept for comparison; sequential, slower on the CPU).
+
+Masks are computed on the host (numpy) from the CPU processor output.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def _blocks(pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> np.ndarray:
+    """(blocks_h, blocks_w, 4*1536) view of the processor output (patches are ordered block by block, 4 per block)."""
+    bh, bw = int(grid_thw[0, 1]) // 2, int(grid_thw[0, 2]) // 2
+    return np.asarray(pixel_values, dtype=np.float32).reshape(bh, bw, -1)
+
+
+def lt_keep_mask(pixel_values: torch.Tensor, grid_thw: torch.Tensor, tau: float) -> np.ndarray:
+    """Keep mask (raster order) of the order-free left-or-upper rule. At least one block is always kept."""
+    blocks = _blocks(pixel_values, grid_thw)
+    drop = np.zeros(blocks.shape[:2], bool)
+    drop[:, 1:] |= np.abs(blocks[:, 1:] - blocks[:, :-1]).max(-1) <= tau
+    drop[1:, :] |= np.abs(blocks[1:] - blocks[:-1]).max(-1) <= tau
+    keep = ~drop
+    if not keep.any():          # a perfectly uniform screen: keep the top-left block so the LLM still sees one visual token
+        keep[0, 0] = True
+    return keep.reshape(-1)
+
+
+def pixelprune_keep_mask(pixel_values: torch.Tensor, grid_thw: torch.Tensor, tau: float) -> np.ndarray:
+    """Keep mask (raster order) of PixelPrune Pred-2D. tau == 0 means bit-exact equality. The first block is always kept.
+
+    Prediction per block (a = left, b = up, c = up-left, all reconstructed):
+        if c == b and a != b: predict a    elif c == a and a != b: predict b    else: predict a
+    First row predicts from the left, first column from above. A dropped block's reconstruction is its prediction.
+    """
+    blocks = _blocks(pixel_values, grid_thw)
+    bh, bw = blocks.shape[:2]
+    rec = blocks.copy()
+    keep = np.ones((bh, bw), bool)
+    if tau == 0:
+        eq = np.array_equal
+    else:
+        def eq(x, y):
+            return float(np.abs(x - y).max()) <= tau
+    for r in range(bh):
+        for c in range(bw):
+            if r == 0 and c == 0:
+                continue
+            if r == 0:
+                pred = rec[0, c - 1]
+            elif c == 0:
+                pred = rec[r - 1, 0]
+            else:
+                a, b, cc = rec[r, c - 1], rec[r - 1, c], rec[r - 1, c - 1]
+                cb, ca, ab = eq(cc, b), eq(cc, a), eq(a, b)
+                pred = a if (cb and not ab) else b if (ca and not ab) else a
+            if eq(blocks[r, c], pred):
+                keep[r, c] = False
+                rec[r, c] = pred
+    return keep.reshape(-1)
+
+
+def _flat_components(eq_left: np.ndarray, eq_up: np.ndarray) -> np.ndarray:
+    """Connected components of the lt equality graph (edge to the left block where eq_left, to the upper block where eq_up).
+    One label per block = the smallest raster index in its component (min-label propagation along the edges + pointer jumping)."""
+    bh, bw = eq_left.shape
+    lab = np.arange(bh * bw).reshape(bh, bw)
+    hl, vu = eq_left[:, 1:], eq_up[1:]
+    while True:
+        old = lab.copy()
+        m = np.minimum(lab[:, 1:], lab[:, :-1])
+        lab[:, 1:] = np.where(hl, m, lab[:, 1:])
+        lab[:, :-1] = np.where(hl, np.minimum(lab[:, :-1], m), lab[:, :-1])
+        m = np.minimum(lab[1:], lab[:-1])
+        lab[1:] = np.where(vu, m, lab[1:])
+        lab[:-1] = np.where(vu, np.minimum(lab[:-1], m), lab[:-1])
+        flat = lab.reshape(-1)
+        lab = flat[flat].reshape(bh, bw)
+        if np.array_equal(lab, old):
+            return lab
+
+
+def gridkeep_keep_mask(pixel_values: torch.Tensor, grid_thw: torch.Tensor, tau: float, tile: int = 2) -> np.ndarray:
+    """Keep mask (raster order) of lt + Grid Keep. The flat components of lt (connected through its left / upper equalities) are
+    cut into tile x tile block tiles; every (component, tile) group without a kept block keeps its raster-first (top-left) block,
+    which is encoded normally at its own position."""
+    blocks = _blocks(pixel_values, grid_thw)
+    bh, bw = blocks.shape[:2]
+    eq_left = np.zeros((bh, bw), bool)
+    eq_up = np.zeros((bh, bw), bool)
+    eq_left[:, 1:] = np.abs(blocks[:, 1:] - blocks[:, :-1]).max(-1) <= tau
+    eq_up[1:, :] = np.abs(blocks[1:] - blocks[:-1]).max(-1) <= tau
+    keep = (~(eq_left | eq_up)).reshape(-1)              # = lt_keep_mask; block (0, 0) is always kept
+    lab = _flat_components(eq_left, eq_up)
+    rr, cc = np.meshgrid(np.arange(bh), np.arange(bw), indexing="ij")
+    ntc = -(-bw // tile)
+    group = (lab * (-(-bh // tile) * ntc) + (rr // tile) * ntc + cc // tile).reshape(-1)   # (component, tile) id
+    orphan = np.flatnonzero(~keep & ~np.isin(group, group[keep]))   # dropped blocks of groups without a kept block, raster order
+    if len(orphan):
+        _, first = np.unique(group[orphan], return_index=True)      # raster-first member of each such group
+        keep[orphan[first]] = True
+    return keep
+
+
+RULES = {"lt": lt_keep_mask, "lt-grid": gridkeep_keep_mask, "pixelprune": pixelprune_keep_mask}
+
+
+def keep_mask(rule: str, pixel_values: torch.Tensor, grid_thw: torch.Tensor, tau: float, grid_tile: int = 3):
+    """Keep mask of `rule`, or None for rule "none" (all blocks kept). `grid_tile` is the Grid Keep tile side (lt-grid only)."""
+    if rule == "none":
+        return None
+    if rule == "lt-grid":
+        return gridkeep_keep_mask(pixel_values, grid_thw, tau, tile=grid_tile)
+    return RULES[rule](pixel_values, grid_thw, tau)

streamembed/engine/temporal.py

@@ -0,0 +1,108 @@
+"""KV Reuse, for temporal redundancy: per-stream cache state, change detection and the refresh rule.
+
+Consecutive screenshots (every 5 s) are often identical; only small regions change (typing, cursor). Across frames, a block counts as
+changed if any pixel differs by more than delta (8 grey levels = 8/255) from when it was LAST computed (not the previous frame, so slow
+drifts cannot accumulate unseen). Changed kept blocks are recomputed; the ViT reuses the cached K/V of all other kept blocks (vit.py),
+and the LLM reuses the cached K/V of their visual tokens (llm.py).
+
+Full refresh (recompute all kept blocks) when the stream starts, when more than rho of the kept blocks changed (a new screen), or every
+K frames since the last refresh (bounds the drift of the approximate reuse).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+
+from .config import MethodConfig
+from .model import Grid, Model
+
+
+def to_device(idx: np.ndarray, m: Model) -> torch.Tensor:
+    # changed: index sets are built on the host and uploaded (MPS index tensors from nonzero() were corrupted under memory pressure)
+    return torch.from_numpy(np.ascontiguousarray(idx, dtype=np.int64)).to(m.device)
+
+
+class StreamState:
+    """Everything one stream carries from frame to frame (one per Engine; reset when a new stream starts)."""
+
+    def __init__(self, m: Model, g: Grid, approx_llm: bool):
+        nl, n, nh, hd = len(m.vit.blocks), g.n_patches, m.vit_heads, m.vit_head_dim
+        out = m.vcfg.out_hidden_size
+        self.key = g.key
+        # ViT: post-RoPE K/V of every patch in every ViT layer, so changed blocks attend over unchanged ones without recomputing them
+        self.vit_k = torch.empty(nl, n, nh, hd, dtype=m.dtype, device=m.device)
+        self.vit_v = torch.empty_like(self.vit_k)
+        # KV Reuse (ViT): per-block ViT outputs (merger = LLM visual token, DeepStack = features added to early LLM layers)
+        self.merged = torch.empty(g.n_blocks, out, dtype=m.dtype, device=m.device)
+        self.deepstack = torch.empty(len(m.deepstack_layers), g.n_blocks, out, dtype=m.dtype, device=m.device)
+        # KV Reuse (ViT): pixels of each block at its last recompute (the change-detection reference)
+        self.ref = torch.empty(g.n_blocks, 4 * m.patch_w.shape[1], dtype=torch.float32, device=m.device)
+        # LLM: post-RoPE K/V of every LLM layer per FULL-grid slot (slot = position in the prompt with all blocks), so a block keeps
+        # its slot even when Token Pruning keeps a different set of blocks next frame (~100 MB fp16)
+        self.llm_k = self.llm_v = None
+        if approx_llm:
+            tc = m.tcfg
+            self.llm_k = torch.empty(tc.num_hidden_layers, g.length, tc.num_key_value_heads, tc.head_dim, dtype=m.dtype, device=m.device)
+            self.llm_v = torch.empty_like(self.llm_k)
+        self.reset(m, g)
+
+    def reset(self, m: Model, g: Grid):
+        self.valid = torch.zeros(g.n_blocks, dtype=torch.bool, device=m.device)   # block has cached ViT state
+        self.since_refresh = 0
+        self.prev_kept: Optional[np.ndarray] = None     # kept blocks of the previous frame (LLM prefix reuse)
+        self.prev_emb: Optional[torch.Tensor] = None    # previous embedding (returned unchanged when nothing changed)
+        self.exact_kv = None                            # exact LLM prefix reuse: HF DynamicCache of the previous frame
+        self.llm_valid = np.zeros(g.length, bool) if self.llm_k is not None else None   # KV Reuse (LLM): slot has cached K/V
+
+
+@dataclass
+class FramePlan:
+    """What to compute for one frame."""
+    kept: np.ndarray                    # block ids fed to ViT/LLM after Token Pruning (sorted)
+    recompute: np.ndarray               # block ids that go through the ViT this frame (subset of kept; = kept on a refresh)
+    refresh: bool
+    changed_frac: float                 # fraction of kept blocks that changed (-1 at stream start / without KV Reuse (ViT))
+    since_refresh: int                  # frames since the last refresh, after this frame
+    state: Optional[StreamState]
+    kept_t: torch.Tensor = None         # device copies of the index sets
+    recompute_t: torch.Tensor = None
+    # LLM part, filled by llm.plan_llm()
+    prefix: int = 0                     # P: positions [0, P) are bit-identical to the previous frame
+    skip_llm: bool = False              # nothing changed at all: reuse the previous embedding
+    llm_rows: Optional[np.ndarray] = None   # KV Reuse (LLM): positions recomputed in the LLM (R)
+
+
+@torch.inference_mode()
+def plan_frame(cfg: MethodConfig, m: Model, g: Grid, state: Optional[StreamState], block_pixels: torch.Tensor, first: bool,
+               keep_mask: Optional[np.ndarray]) -> FramePlan:
+    """Kept blocks (Token Pruning) and recomputed blocks (KV Reuse) of one frame. block_pixels: (n_blocks, 4*1536) float32 on the
+    device."""
+    kept = np.flatnonzero(keep_mask) if keep_mask is not None else np.arange(g.n_blocks)
+    kept_t = to_device(kept, m)
+    if state is None:                                   # no KV Reuse (ViT): every kept block goes through the ViT
+        return FramePlan(kept, kept, True, -1.0, 0, None, kept_t, kept_t)
+    refresh = first or state.since_refresh + 1 >= cfg.refresh_every
+    frac = -1.0
+    changed = None
+    if not first:
+        changed = ((block_pixels - state.ref).abs().amax(1) > cfg.delta) | ~state.valid
+        frac = float(changed[kept_t].float().mean())    # float32 mean on the device (rounding decides exact ties at rho)
+        refresh = refresh or frac > cfg.refresh_rho
+    if refresh:
+        recompute, recompute_t = kept, kept_t
+    else:
+        recompute = kept[changed.cpu().numpy()[kept]]
+        recompute_t = to_device(recompute, m)
+    return FramePlan(kept, recompute, refresh, frac, 0 if refresh else state.since_refresh + 1, state, kept_t, recompute_t)
+
+
+def commit(plan: FramePlan, emb: torch.Tensor):
+    """After the frame: remember what the next frame's prefix reuse / skip needs."""
+    st = plan.state
+    if st is not None:
+        st.prev_kept = plan.kept
+        st.prev_emb = emb
+        st.since_refresh = plan.since_refresh

streamembed/engine/vit.py

@@ -0,0 +1,86 @@
+"""ViT forward over a subset of blocks: plain (full screenshot / Token Pruning) or partial with cached K/V (KV Reuse (ViT)).
+
+Re-implements the HF Qwen3-VL vision forward layer by layer (same modules and weights) so that only some patches are queries:
+  * Token Pruning: only the patches of kept blocks enter (changed: stock HF always runs every patch).
+  * KV Reuse: only patches of recomputed blocks are queries in EVERY layer; their fresh post-RoPE K/V are written into the stream's cache
+    and they attend over the cached K/V of all kept blocks. Unchanged blocks are never run; their merger / DeepStack outputs come from a
+    per-block cache. Attention in the Qwen3-VL ViT is full (non-causal) over the image, so unchanged blocks do not see the change: this is
+    the approximation of KV Reuse (ViT).
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers.models.qwen3_vl import modeling_qwen3_vl as MQ
+
+from .model import Grid, Model, embed_patches
+from .temporal import FramePlan, to_device
+
+
+def patch_rows(blocks: np.ndarray) -> np.ndarray:
+    """Patch row ids of blocks (each block = 4 consecutive patch rows in the processor's order)."""
+    return (blocks[:, None] * 4 + np.arange(4)).reshape(-1)
+
+
+@torch.inference_mode()
+def run_vit(m: Model, g: Grid, plan: FramePlan, pixels: torch.Tensor, block_pixels: torch.Tensor, linear_patch_embed: bool):
+    """Returns (visual tokens of the kept blocks, list of DeepStack features of the kept blocks), both in kept order."""
+    st = plan.state
+    if len(plan.recompute) == 0:                        # nothing changed: everything from the cache
+        return _from_cache(m, plan)
+    V = m.vit
+    q_rows = to_device(patch_rows(plan.recompute), m)
+    partial = st is not None and len(plan.recompute) < len(plan.kept)
+    kv_rows = to_device(patch_rows(plan.kept), m) if partial else None
+    x = embed_patches(m, pixels[q_rows].to(m.dtype), linear_patch_embed)
+    x = x + g.vit_pos[q_rows].to(x.dtype)
+    cos, sin = g.vit_cos[q_rows], g.vit_sin[q_rows]
+    n = x.shape[0]
+    deepstack = []
+    for layer, blk in enumerate(V.blocks):
+        q, k, v = KERNELS["vit_qkv"](blk, x, cos, sin, m.vit_heads, m.vit_head_dim)
+        if st is not None:                              # changed: keep fresh K/V for later frames (also on a refresh)
+            st.vit_k[layer][q_rows] = k
+            st.vit_v[layer][q_rows] = v
+            if partial:                                 # changed: attend over all kept blocks, unchanged ones from the cache
+                k, v = st.vit_k[layer][kv_rows], st.vit_v[layer][kv_rows]
+        # changed: one SDPA call over a single image (no cu_seqlens varlen path); identical for batch-1 inputs
+        o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], k.transpose(0, 1)[None], v.transpose(0, 1)[None],
+                                           dropout_p=0.0, scale=blk.attn.scaling, is_causal=False)
+        x = KERNELS["vit_out"](blk, x, o, n)
+        if layer in m.deepstack_layers:
+            deepstack.append(V.deepstack_merger_list[m.deepstack_layers.index(layer)](x))
+    merged = V.merger(x)
+    if st is None:
+        return merged, deepstack
+    c = plan.recompute_t
+    st.merged[c] = merged
+    for i in range(len(deepstack)):
+        st.deepstack[i][c] = deepstack[i]
+    st.valid[c] = True
+    st.ref[c] = block_pixels[c]
+    return _from_cache(m, plan)
+
+
+def _vit_qkv(blk, x, cos, sin, heads: int, head_dim: int):
+    """norm1 -> qkv projection -> 2-D RoPE (the dense part before attention)."""
+    h = blk.norm1(x)
+    q, k, v = blk.attn.qkv(h).reshape(-1, 3, heads, head_dim).permute(1, 0, 2, 3).unbind(0)
+    q, k = MQ.apply_rotary_pos_emb_vision(q, k, cos, sin)
+    return q, k, v
+
+
+def _vit_out(blk, x, o, n: int):
+    """attention output projection + residual, then norm2 -> MLP + residual (the dense part after attention)."""
+    x = x + blk.attn.proj(o[0].transpose(0, 1).reshape(n, -1))
+    return x + blk.mlp(blk.norm2(x))
+
+
+# changed: the dense per-block parts are looked up here so cpu_compile.enable() can swap in torch.compile'd versions (same ops, same order)
+KERNELS = {"vit_qkv": _vit_qkv, "vit_out": _vit_out}
+
+
+def _from_cache(m: Model, plan: FramePlan):
+    st, kept = plan.state, plan.kept_t
+    return st.merged[kept], [st.deepstack[i][kept] for i in range(len(m.deepstack_layers))]

streamembed/gate.py

@@ -0,0 +1,102 @@
+"""The "if changed" capture gate of the paper's capture-interval experiment (policy G5_0.1).
+
+A captured screenshot is encoded if it is the first one, or if more than 10 % of its Grid keep 3x3 kept blocks (32x32 RGB) differ by more
+than 32/255 from the last encoded frame, or as a heartbeat 30 s after the last encode.
+
+Numpy port of `timing/c2_gate.py`, generalised from a fixed 1280x704 grid to any resolution divisible by 32. One addition to the paper's
+gate: a heartbeat is skipped when no block differs by more than 8/255 at all. An idle screen would otherwise add identical rows every 30 s,
+and near-duplicates crowd the top results.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from PIL import Image
+
+BLOCK = 32
+TAU = 8             # a block is flat if it matches its left or upper neighbour within this
+DIFF = 32           # a block changed if any value differs from the last encoded frame by more than this
+FRACTION = 0.1      # encode if more than this fraction of the kept blocks changed
+HEARTBEAT_S = 30.0
+TILE = 3
+IDLE_DIFF = 8       # below this everywhere, the screen did not change
+
+
+def blocks(img: Image.Image) -> np.ndarray:
+    """(blocks_h, blocks_w, 3 * 32 * 32) uint8 array of an RGB image whose sides are multiples of 32."""
+    a = np.asarray(img.convert("RGB"))
+    h, w = a.shape[0] // BLOCK, a.shape[1] // BLOCK
+    assert a.shape[0] == h * BLOCK and a.shape[1] == w * BLOCK, f"sides must be multiples of {BLOCK}: {a.shape}"
+    return a.reshape(h, BLOCK, w, BLOCK, 3).transpose(0, 2, 1, 3, 4).reshape(h, w, -1)
+
+
+def _max_abs_diff(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """Per-block max |a - b| on the last axis, without widening the uint8 values."""
+    return (np.maximum(a, b) - np.minimum(a, b)).max(-1)
+
+
+def _components(eq_left: np.ndarray, eq_up: np.ndarray) -> np.ndarray:
+    """Connected components of the flat-block graph (min-label propagation, with pointer jumping)."""
+    h, w = eq_left.shape
+    lab = np.arange(h * w).reshape(h, w)
+    while True:
+        old = lab.copy()
+        m = np.minimum(lab[:, 1:], lab[:, :-1])
+        lab[:, 1:] = np.where(eq_left[:, 1:], m, lab[:, 1:])
+        lab[:, :-1] = np.where(eq_left[:, 1:], np.minimum(lab[:, :-1], m), lab[:, :-1])
+        m = np.minimum(lab[1:], lab[:-1])
+        lab[1:] = np.where(eq_up[1:], m, lab[1:])
+        lab[:-1] = np.where(eq_up[1:], np.minimum(lab[:-1], m), lab[:-1])
+        lab = lab.reshape(-1)
+        lab = lab[lab].reshape(h, w)
+        if np.array_equal(lab, old):
+            return lab
+
+
+def kept_blocks(b: np.ndarray) -> np.ndarray:
+    """Block ids (raster order) kept by spatial pruning with Grid keep: non-flat blocks, plus one block per (flat component, 3x3 tile)."""
+    h, w = b.shape[:2]
+    eq_left, eq_up = np.zeros((h, w), bool), np.zeros((h, w), bool)
+    eq_left[:, 1:] = _max_abs_diff(b[:, 1:], b[:, :-1]) <= TAU
+    eq_up[1:, :] = _max_abs_diff(b[1:], b[:-1]) <= TAU
+    keep = ~(eq_left | eq_up).reshape(-1)
+    if not keep.any():
+        keep[0] = True
+    rows, cols = np.meshgrid(np.arange(h), np.arange(w), indexing="ij")
+    tiles_per_row = -(-w // TILE)
+    tile = ((rows // TILE) * tiles_per_row + cols // TILE).reshape(-1)
+    key = _components(eq_left, eq_up).reshape(-1) * (-(-h // TILE) * tiles_per_row) + tile
+    covered = np.isin(key, np.unique(key[keep]))
+    todo = np.flatnonzero(~keep & ~covered)
+    if len(todo):
+        uniq, inv = np.unique(key[todo], return_inverse=True)
+        first = np.full(len(uniq), -1)
+        first[inv[::-1]] = todo[::-1]
+        keep[first] = True
+    return np.flatnonzero(keep)
+
+
+@dataclass
+class Decision:
+    encode: bool
+    idle: bool          # nothing on screen changed since the last encoded frame
+
+
+class Gate:
+    def __init__(self):
+        self.ref: np.ndarray | None = None
+        self.last_ts = 0.0
+
+    def __call__(self, img: Image.Image, ts: float) -> Decision:
+        b = blocks(img)
+        if self.ref is None or self.ref.shape != b.shape:
+            encode, idle = True, False
+        else:
+            diff = _max_abs_diff(b, self.ref).reshape(-1)
+            idle = bool((diff <= IDLE_DIFF).all())
+            heartbeat = ts - self.last_ts >= HEARTBEAT_S
+            encode = not idle and (heartbeat or float((diff > DIFF)[kept_blocks(b)].mean()) > FRACTION)
+        if encode:
+            self.ref, self.last_ts = b, ts
+        return Decision(encode, idle)

streamembed/index.py

@@ -0,0 +1,138 @@
+"""The index folder: one fp16 vector and one JSON line per stored frame, appended in stream order.
+
+    index.json      encoder name; vectors from different encoders must not share an index
+    vectors.f16     2048 fp16 per frame; row i belongs to line i of frames.jsonl
+    frames.jsonl    {"ts", "path", "app", "title"} per frame
+    seen.jsonl      [row, ts]: the screen of that row was still on screen at ts (unchanged captures)
+    images/         the stored JPEGs
+    mu.npy, codes.bin   "binary" only: centring vector and 64-B sign codes of the first 512 dimensions
+
+Search is exact ("flat"), or for "binary" above BINARY_FROM frames: Hamming top 1,000 over the codes, then exact re-rank (the paper's index).
+"""
+from __future__ import annotations
+
+import json
+import threading
+from pathlib import Path
+
+import numpy as np
+
+DIM, CODE_DIM, CANDIDATES = 2048, 512, 1000
+BINARY_FROM = 1000
+INDEX_TYPES = ("flat", "binary")
+
+
+def _codes(v: np.ndarray, mu: np.ndarray) -> np.ndarray:
+    t = v[:, :CODE_DIM].astype(np.float32)
+    t /= np.linalg.norm(t, axis=1, keepdims=True).clip(1e-9)
+    return np.packbits(t >= mu, axis=1)
+
+
+class Index:
+    def __init__(self, root: Path, index_type: str, encoder_name: str):
+        if index_type not in INDEX_TYPES:
+            raise ValueError(f"Invalid index type: {index_type} not in {list(INDEX_TYPES)}")
+        self.root, self.type, self.lock = Path(root), index_type, threading.RLock()
+        self.images = self.root / "images"
+        self.images.mkdir(parents=True, exist_ok=True)
+        self.root.chmod(0o700)
+        info = self.root / "index.json"
+        if info.exists() and json.loads(info.read_text())["encoder"] != encoder_name:
+            raise ValueError(f"{root} was built with encoder {json.loads(info.read_text())['encoder']!r}, not {encoder_name!r}. "
+                             "Use another index_dir.")
+        info.write_text(json.dumps({"encoder": encoder_name}))
+        self.vec_path, self.meta_path, self.seen_path = self.root / "vectors.f16", self.root / "frames.jsonl", self.root / "seen.jsonl"
+        for p in (self.vec_path, self.meta_path, self.seen_path):
+            p.touch()
+        self.frames = [json.loads(l) for l in self.meta_path.read_text().splitlines() if l.strip()]
+        for m in self.frames:
+            m["ts_last"] = m["ts"]
+        n = min(len(self.frames), self.vec_path.stat().st_size // (DIM * 2))      # a crash between the two appends leaves one longer
+        self._truncate(n)
+        for row, ts in (json.loads(l) for l in self.seen_path.read_text().splitlines() if l.strip()):
+            if row < n:
+                self.frames[row]["ts_last"] = max(self.frames[row]["ts_last"], ts)
+        self.mu = self.codes = None
+        if self.type == "binary" and n > BINARY_FROM:
+            self._load_codes()
+
+    def __len__(self) -> int:
+        return len(self.frames)
+
+    def _truncate(self, n: int) -> None:
+        self.frames = self.frames[:n]
+        with open(self.vec_path, "r+b") as f:
+            f.truncate(n * DIM * 2)
+        self.meta_path.write_text("".join(json.dumps({k: m[k] for k in ("ts", "path", "app", "title")}) + "\n" for m in self.frames))
+
+    def _rows(self) -> np.ndarray:
+        return np.memmap(self.vec_path, np.float16, "r", shape=(len(self.frames), DIM))
+
+    def _load_codes(self) -> None:
+        mu_path, code_path = self.root / "mu.npy", self.root / "codes.bin"
+        if not mu_path.exists():
+            sample = self._rows()[np.linspace(0, len(self) - 1, min(len(self), 50_000), dtype=np.int64)]
+            t = np.asarray(sample[:, :CODE_DIM], np.float32)
+            np.save(mu_path, (t / np.linalg.norm(t, axis=1, keepdims=True).clip(1e-9)).mean(0))
+            code_path.write_bytes(b"")
+        self.mu = np.load(mu_path)
+        self.codes = np.fromfile(code_path, np.uint8).reshape(-1, CODE_DIM // 8)[:len(self)]
+        if len(self.codes) < len(self):                                          # codes lag after a crash or a "flat" session
+            self.codes = np.concatenate([self.codes, _codes(np.asarray(self._rows()[len(self.codes):]), self.mu)])
+        self.codes.tofile(code_path)
+
+    def add(self, vector: np.ndarray, meta: dict) -> int:
+        with self.lock:
+            v = vector.astype(np.float16).reshape(1, DIM)
+            with open(self.vec_path, "ab") as f:          # vector first: on reopen, a vector without a line is dropped
+                v.tofile(f)
+            line = {k: meta[k] for k in ("ts", "path", "app", "title")}
+            with open(self.meta_path, "a") as f:
+                f.write(json.dumps(line) + "\n")
+            self.frames.append({**line, "ts_last": line["ts"]})
+            if self.codes is not None:
+                c = _codes(v, self.mu)
+                self.codes = np.concatenate([self.codes, c])
+                with open(self.root / "codes.bin", "ab") as f:
+                    c.tofile(f)
+            elif self.type == "binary" and len(self) > BINARY_FROM:
+                self._load_codes()
+            return len(self) - 1
+
+    def still_on_screen(self, ts: float) -> None:
+        """The last stored screen is unchanged at ts."""
+        with self.lock:
+            if self.frames:
+                self.frames[-1]["ts_last"] = ts
+                with open(self.seen_path, "a") as f:
+                    f.write(json.dumps([len(self) - 1, ts]) + "\n")
+
+    def search(self, q: np.ndarray, k: int, dedupe_s: float = 30.0) -> list[tuple[int, float]]:
+        """(row, cosine) best first. A match within dedupe_s seconds of a better one is dropped (screens stay up for a while)."""
+        with self.lock:
+            n = len(self)
+            if n == 0:
+                return []
+            rows, q = self._rows(), q.astype(np.float32)
+            if self.codes is None:
+                ids = np.arange(n)
+                scores = np.concatenate([np.asarray(rows[s:s + 8192], np.float32) @ q for s in range(0, n, 8192)])
+            else:
+                t = q[:CODE_DIM] / max(float(np.linalg.norm(q[:CODE_DIM])), 1e-9)
+                qc = np.packbits(t >= self.mu).view(np.uint64)
+                dist = np.bitwise_count(self.codes.view(np.uint64) ^ qc).sum(1)
+                ids = np.sort(np.argpartition(dist, min(CANDIDATES, n) - 1)[:CANDIDATES])
+                scores = np.asarray(rows[ids], np.float32) @ q
+            out, kept = [], []
+            for j in np.argsort(-scores):
+                ts = self.frames[ids[j]]["ts"]
+                if all(abs(ts - t) >= dedupe_s for t in kept):
+                    kept.append(ts)
+                    out.append((int(ids[j]), float(scores[j])))
+                    if len(out) == k:
+                        break
+            return out
+
+    def around(self, row: int, n: int) -> list[int]:
+        """Rows recorded just before and after `row`, including it (stream order is time order)."""
+        return list(range(max(0, row - n), min(len(self), row + n + 1)))

streamembed/models.py

@@ -0,0 +1,90 @@
+"""Which model encodes screens and questions. A model is a base model from its own Hugging Face repo, pinned to a commit, plus two
+optional LoRAs (frame adapter, question LoRA) from the StreamEmbed model repo. `models.json` in that repo lists them:
+
+    {"qwen3-vl-2b": {"base": "Qwen/Qwen3-VL-Embedding-2B", "revision": "<commit>",
+                     "frame_lora": "qwen3-vl-2b/frame_adapter.safetensors", "query_lora": "qwen3-vl-2b/query_lora.safetensors",
+                     "scale": 1.0, "architecture": "qwen3-vl", "description": "..."}}
+
+The model repo is a Hugging Face repo id, or a local folder with the same layout (set by `configure(models_repo=...)` or
+$STREAMEMBED_MODELS). Base models without LoRAs are built in, so they work without the model repo.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+MODELS_REPO = os.environ.get("STREAMEMBED_MODELS", "dayoon/streamembed")
+DEFAULT_MODEL = "qwen3-vl-2b"
+STREAM_ARCHITECTURES = ("qwen3-vl",)      # the stream encoder re-implements the Qwen3-VL forward pass
+
+BUILTIN = {
+    "qwen3-vl-2b-base": dict(base="Qwen/Qwen3-VL-Embedding-2B", revision="9f2f7e710d6d81056aa5c0a4f04764fec6bb7bda", architecture="qwen3-vl",
+                             description="Qwen3-VL-Embedding-2B without LoRAs"),
+}
+
+
+@dataclass
+class ModelSpec:
+    name: str
+    base: str
+    revision: str | None
+    frame_lora: Path | None
+    query_lora: Path | None
+    scale: float
+    architecture: str
+
+    @property
+    def weights(self) -> tuple[Path, Path] | None:
+        return (self.frame_lora, self.query_lora) if self.frame_lora and self.query_lora else None
+
+    @property
+    def encoder_name(self) -> str:
+        """Recorded in each index: vectors from another model or base commit must not share it."""
+        return f"{self.name}@{(self.revision or 'main')[:12]}"
+
+
+def _fetch(repo: str, filename: str) -> Path:
+    if Path(repo).expanduser().is_dir():
+        p = Path(repo).expanduser() / filename
+        if not p.is_file():
+            raise FileNotFoundError(f"{p} not found")
+        return p
+    from huggingface_hub import hf_hub_download
+    try:
+        return Path(hf_hub_download(repo, filename, local_files_only=True))     # once cached, no network
+    except Exception:
+        return Path(hf_hub_download(repo, filename))
+
+
+def _manifest(repo: str) -> dict:
+    return json.loads(_fetch(repo, "models.json").read_text())
+
+
+def available_models(repo: str | None = None) -> list[str]:
+    """Model names: the model repo's, then the built-in base models. The repo's are left out if it cannot be reached."""
+    try:
+        names = list(_manifest(repo or MODELS_REPO))
+    except Exception:
+        names = []
+    return names + [n for n in BUILTIN if n not in names]
+
+
+def resolve(name: str, repo: str | None = None) -> ModelSpec:
+    repo = repo or MODELS_REPO
+    if name in BUILTIN:
+        entry = BUILTIN[name]
+    else:
+        try:
+            listing = _manifest(repo)
+        except Exception as e:
+            raise RuntimeError(f"Could not read models.json from {repo!r} ({e}). Built-in models: {list(BUILTIN)}") from e
+        if name not in listing:
+            raise ValueError(f"Unknown model {name!r}; available: {list(listing) + list(BUILTIN)}")
+        entry = listing[name]
+    arch = entry.get("architecture", "qwen3-vl")
+    if arch not in STREAM_ARCHITECTURES:
+        raise NotImplementedError(f"{name}: architecture {arch!r} is not supported yet; the stream encoder supports {STREAM_ARCHITECTURES}")
+    loras = [_fetch(repo, entry[k]) if entry.get(k) else None for k in ("frame_lora", "query_lora")]
+    return ModelSpec(name, entry["base"], entry.get("revision"), loras[0], loras[1], float(entry.get("scale", 1.0)), arch)