Take the dataset repo as user input and stream windows on demand

The repo is now typed into the app rather than baked in via DATASET_REPO,
which becomes only the prefilled default. Accepts owner/name, a revision
suffix, or a pasted dataset URL.

Generalises the layout assumptions so an arbitrary repo renders: the manifest
and metrics files are optional, the window list falls back to listing
windows/ one level deep, per-window metadata is backfilled from meta.json,
the camera count is inferred from the frame aspect ratio, and window length
and critical-frame markers are no longer fixed at 33/16.

Storage is now per session and bounded, so large datasets are never pulled
whole: loading a repo fetches only the index, window files arrive when the
window is opened, at most MAX_WINDOWS_ON_DISK stay resident, and the
session's directory is deleted when the tab closes. A janitor reaps idle
sessions and directories orphaned by an earlier process.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

README.md

@@ -13,51 +13,73 @@ short_description: Frame-locked viewer for rollouts vs. ground truth
 
 # Rollout vs. ground truth
 
-Side-by-side viewer for forward-dynamics world-model rollouts. Reads windows straight from a
-Hugging Face **dataset** repo — the Space holds no data of its own.
+Side-by-side viewer for forward-dynamics world-model rollouts. Type a **dataset repo id** into
+the box at the top and it renders that repo — the Space holds no data of its own and is not tied
+to any one dataset.
 
-Default dataset: [`OneWorld-AI/abc-critical-window-rollout`](https://huggingface.co/datasets/OneWorld-AI/abc-critical-window-rollout)
-(300 held-out ABC bimanual-YAM windows, each a 33-frame slice centred on a gripper critical frame).
+Accepted forms: `owner/name`, `owner/name@revision`, or a pasted `https://huggingface.co/datasets/…`
+URL. `DATASET_REPO` sets what is prefilled at startup.
 
 ## What it shows
 
 - **Playback** — GT and generated rendered into one video, so the two panels stay frame-locked
   instead of drifting apart the way two independent players do. Optional `|GT − generated|`
-  difference panel. Playback fps is adjustable (the 33-frame window is only ~1 s at native 30 fps).
-- **Frame stepper** — a single frame from both videos side by side, defaulting to frame 16
-  (the critical frame).
-- **Driving actions** — the `[33, 14]` action chunk fed to the model, split into grippers /
-  left-arm joints / right-arm joints, with the critical frame marked.
-- **All windows** — sortable table of every window with its PSNR; click a row to load it.
+  difference panel. Playback fps is adjustable (a 33-frame window is only ~1 s at 30 fps).
+- **Frame stepper** — a single frame from both videos side by side, defaulting to the critical
+  frame when the dataset marks one.
+- **Driving actions** — the action chunk fed to the model. A 14-D chunk is split into grippers /
+  left-arm / right-arm; any other width is plotted as raw channels.
+- **All windows** — sortable table of every window; click a row to load it.
 
-Filter by task, gripper edge (grasp / release) and arm; sort by critical-frame PSNR to jump
-straight to the hardest or easiest cases.
+Filter by task, gripper edge and arm; sort by critical-frame PSNR to jump to the hardest or
+easiest cases.
 
 ## Expected dataset layout
 
 ```
-windows_manifest.json                 # list of window dicts: episode_id, task, edge_type, arm,
-                                      #   start_frame, critical_frame, amplitude, n_frames, ...
-metrics_summary.json                  # list of {id, task, edge, arm, full_mse, full_psnr,
-                                      #   crit_mse, crit_psnr}
-windows/<episode_id>/gt.mp4           # 33 frames, 3 cameras vstacked (top / L wrist / R wrist)
-windows/<episode_id>/generated.mp4    # model rollout for the same slice, same geometry
-windows/<episode_id>/actions.json     # [33, 14] = [L arm 6, L grip, R arm 6, R grip]
-windows/<episode_id>/meta.json
+windows_manifest.json                 # optional: list of window dicts (episode_id, task,
+                                      #   edge_type, arm, start_frame, critical_frame, ...)
+metrics_summary.json                  # optional: [{id, task, edge, arm, full_mse, full_psnr,
+                                      #   crit_mse, crit_psnr}, ...]
+windows/<episode_id>/gt.mp4           # required
+windows/<episode_id>/generated.mp4    # required, same geometry and length as gt.mp4
+windows/<episode_id>/actions.json     # optional: [T, D] driving actions
+windows/<episode_id>/meta.json        # optional: per-window metadata
 ```
 
-Any dataset matching this layout works. Black letterbox bars are trimmed automatically, and the
-per-camera split is derived from the frame height, so other stacked-camera resolutions are fine.
+Only the two mp4s are required. With no `windows_manifest.json` the window list is discovered by
+listing `windows/` one level deep, and per-window metadata is read from each `meta.json` as you
+open it. Videos may stack **N cameras vertically** — N is inferred from the frame aspect ratio,
+and the camera selector adapts. Black letterbox bars are trimmed automatically, identically for
+GT and generated.
 
-## Configuration
+## Storage model
+
+Nothing is downloaded at build time, and no repo is ever pulled whole — important for datasets
+too large to fit on the Space.
+
+- Loading a repo fetches only the manifest and metrics JSONs (a few hundred KB), or one directory
+  listing when there is no manifest.
+- A window's mp4s, `actions.json` and `meta.json` are fetched **the first time that window is
+  opened**, never in bulk.
+- Each browser session gets its own scratch directory. At most `MAX_WINDOWS_ON_DISK` windows stay
+  resident; beyond that the least recently used are deleted. Decoded frames are capped at 2
+  windows in RAM and composed strips at 1.
+- The whole directory is deleted when the session ends. A janitor also reaps sessions idle past
+  `SESSION_TTL` and directories orphaned by an earlier process, so nothing accumulates if a tab
+  dies without a clean shutdown.
 
-| Variable | Purpose |
-|---|---|
-| `DATASET_REPO` | Dataset repo to read (default `OneWorld-AI/abc-critical-window-rollout`). |
-| `HF_TOKEN` | **Required if the dataset is private.** Add it under *Settings → Variables and secrets → Secrets*. |
+The sidebar shows live occupancy (`cached this session: n/N windows`).
+
+## Configuration
 
-Only `windows_manifest.json` and `metrics_summary.json` are fetched at startup; videos and actions
-are downloaded per window on demand and cached, as are the composed comparison videos.
+| Variable | Default | Purpose |
+|---|---|---|
+| `DATASET_REPO` | `OneWorld-AI/abc-critical-window-rollout` | Repo prefilled in the input box. |
+| `HF_TOKEN` | — | **Required for private datasets.** Set under *Settings → Variables and secrets*. |
+| `MAX_WINDOWS_ON_DISK` | `12` | Windows kept per session before LRU eviction. |
+| `MAX_DISCOVER` | `5000` | Cap on windows listed when a repo has no manifest. |
+| `SESSION_TTL` | `3600` | Seconds of idleness before a session's storage is reclaimed. |
 
 ## Running locally
 

app.py

@@ -1,47 +1,68 @@
 """
-Side-by-side viewer for Cosmos forward-dynamics rollouts vs. ground truth.
+Side-by-side viewer for forward-dynamics rollouts vs. ground truth.
 
-Data lives in a separate HF dataset repo laid out as:
+Point it at any HF dataset repo laid out as:
 
-    windows_manifest.json          # 300 window definitions
-    metrics_summary.json           # per-window GT-vs-generated MSE / PSNR
-    windows/<episode_id>/gt.mp4         # 33 frames, 224x672 (3 cams vstacked)
-    windows/<episode_id>/generated.mp4  # same slice, model rollout
-    windows/<episode_id>/actions.json   # [33, 14] driving action chunk
-    windows/<episode_id>/meta.json
+    windows_manifest.json               # optional: list of window dicts
+    metrics_summary.json                # optional: per-window MSE / PSNR
+    windows/<episode_id>/gt.mp4         # ground-truth window
+    windows/<episode_id>/generated.mp4  # model rollout for the same slice
+    windows/<episode_id>/actions.json   # optional: [T, D] driving actions
+    windows/<episode_id>/meta.json      # optional: per-window metadata
 
-Set DATASET_REPO to point at a different dataset. If the dataset is private,
-the Space needs an HF_TOKEN secret with read access to it.
+Only `windows/<id>/{gt,generated}.mp4` is required. Without a manifest the
+window list is discovered by listing `windows/` one level deep. Videos may
+stack N cameras vertically; N is inferred from the frame aspect ratio.
+
+Nothing is downloaded at build time and no repo is ever pulled whole. Each
+browser session gets its own scratch directory; window files arrive only when
+that window is opened, at most MAX_WINDOWS_ON_DISK are kept, and the whole
+directory is deleted when the session ends.
+
+DATASET_REPO sets the repo prefilled at startup. Private repos need an
+HF_TOKEN secret with read access.
 """
 
 from __future__ import annotations
 
-import functools
 import json
 import os
+import re
+import shutil
 import tempfile
+import threading
+import time
+import uuid
+from collections import OrderedDict
 from pathlib import Path
 
 import cv2
 import gradio as gr
 import imageio_ffmpeg
 import numpy as np
-from huggingface_hub import hf_hub_download
+from huggingface_hub import HfApi, RepoFolder, hf_hub_download
 from matplotlib.figure import Figure
 
-DATASET_REPO = os.environ.get("DATASET_REPO", "OneWorld-AI/abc-critical-window-rollout")
+DEFAULT_REPO = os.environ.get("DATASET_REPO", "OneWorld-AI/abc-critical-window-rollout")
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None
 
-CACHE_DIR = Path(tempfile.gettempdir()) / "wmviz_render"
-CACHE_DIR.mkdir(parents=True, exist_ok=True)
+BASE_DIR = Path(tempfile.gettempdir()) / "wmviz"
+BASE_DIR.mkdir(parents=True, exist_ok=True)
 
-# Each video stacks three 224x224 camera feeds vertically.
-VIEWS = ["All cameras", "Top", "Left wrist", "Right wrist"]
-VIEW_ROW = {"Top": 0, "Left wrist": 1, "Right wrist": 2}
+MAX_WINDOWS_ON_DISK = int(os.environ.get("MAX_WINDOWS_ON_DISK", 12))
+MAX_PAIRS_IN_RAM = 2
+MAX_COMPOSED_IN_RAM = 1
+MAX_DISCOVER = int(os.environ.get("MAX_DISCOVER", 5000))
+SESSION_TTL = int(os.environ.get("SESSION_TTL", 3600))
+
+# Conventional names when a video stacks exactly three cameras.
+THREE_CAM_NAMES = ["Top", "Left wrist", "Right wrist"]
+ALL_CAMS = -1
+NO_CRIT = -1
 
 FONT = cv2.FONT_HERSHEY_SIMPLEX
 HDR, FTR, SEP = 26, 26, 4
-SCALE = 2  # native frames are only 224 wide; upscale so overlays stay legible
+SCALE = 2  # source frames are small; upscale so overlays stay legible
 
 GT_COLOR = (130, 225, 140)
 GEN_COLOR = (120, 180, 255)
@@ -56,62 +77,281 @@ SORTS = [
 ]
 
 
-# --------------------------------------------------------------------------- data
-
-
-def fetch(rel_path: str) -> str:
-    """Download one file from the dataset repo (hf_hub_download caches on disk)."""
-    return hf_hub_download(
-        repo_id=DATASET_REPO,
-        filename=rel_path,
+# --------------------------------------------------------------------------- session store
+
+
+def parse_repo(text: str) -> tuple[str, str | None]:
+    """Split "owner/name@revision" into (repo_id, revision). Tolerates a pasted URL."""
+    text = (text or "").strip().rstrip("/")
+    m = re.match(r"^https?://huggingface\.co/(?:datasets/)?(.+?)(?:/tree/([^/]+))?$", text)
+    if m:
+        return m.group(1), m.group(2)
+    if "@" in text:
+        repo, _, rev = text.partition("@")
+        return repo.strip(), rev.strip() or None
+    return text, None
+
+
+class Session:
+    """One browser session's scratch space: bounded on disk and in RAM."""
+
+    def __init__(self, sid: str):
+        self.sid = sid
+        self.root = BASE_DIR / sid
+        self.data = self.root / "data"
+        self.render = self.root / "render"
+        for d in (self.data, self.render):
+            d.mkdir(parents=True, exist_ok=True)
+        self.repo = ""
+        self.revision: str | None = None
+        self.rows: list[dict] = []
+        self.note = ""
+        self._on_disk: OrderedDict[str, bool] = OrderedDict()
+        self._pairs: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict()
+        self._composed: OrderedDict[tuple, np.ndarray] = OrderedDict()
+        self._meta: dict[str, dict] = {}
+        self.touched = time.time()
+
+    # -- fetching ----------------------------------------------------------
+
+    def fetch(self, rel: str, optional: bool = False) -> str | None:
+        """Download one repo file into this session's directory."""
+        self.touched = time.time()
+        try:
+            return hf_hub_download(
+                repo_id=self.repo,
+                filename=rel,
+                repo_type="dataset",
+                revision=self.revision,
+                token=HF_TOKEN,
+                local_dir=str(self.data),
+            )
+        except Exception:
+            if optional:
+                return None
+            raise
+
+    def window_file(self, episode_id: str, name: str, optional: bool = False) -> str | None:
+        path = self.fetch(f"windows/{episode_id}/{name}", optional=optional)
+        if path:
+            self._keep(episode_id)
+        return path
+
+    def _keep(self, episode_id: str) -> None:
+        """Register a window as resident and evict the least recently used ones."""
+        self._on_disk[episode_id] = True
+        self._on_disk.move_to_end(episode_id)
+        while len(self._on_disk) > MAX_WINDOWS_ON_DISK:
+            old, _ = self._on_disk.popitem(last=False)
+            shutil.rmtree(self.data / "windows" / old, ignore_errors=True)
+            self._pairs.pop(old, None)
+            for key in [k for k in self._composed if k[0] == old]:
+                self._composed.pop(key, None)
+            for f in self.render.glob(f"{old}_*.mp4"):
+                f.unlink(missing_ok=True)
+
+    def disk_windows(self) -> int:
+        return len(self._on_disk)
+
+    # -- bounded RAM caches ------------------------------------------------
+
+    def pair(self, episode_id: str) -> tuple[np.ndarray, np.ndarray]:
+        hit = self._pairs.get(episode_id)
+        if hit is not None:
+            self._pairs.move_to_end(episode_id)
+            return hit
+        gt = read_video(self.window_file(episode_id, "gt.mp4"))
+        gen = read_video(self.window_file(episode_id, "generated.mp4"))
+        t = min(len(gt), len(gen))
+        value = (gt[:t], gen[:t])
+        self._pairs[episode_id] = value
+        while len(self._pairs) > MAX_PAIRS_IN_RAM:
+            self._pairs.popitem(last=False)
+        return value
+
+    def composed(self, episode_id: str, view: int, show_diff: bool, crit: int) -> np.ndarray:
+        key = (episode_id, view, show_diff, crit)
+        hit = self._composed.get(key)
+        if hit is not None:
+            self._composed.move_to_end(key)
+            return hit
+        value = compose(self.pair(episode_id), view, show_diff, crit)
+        self._composed[key] = value
+        while len(self._composed) > MAX_COMPOSED_IN_RAM:
+            self._composed.popitem(last=False)
+        return value
+
+    def meta(self, episode_id: str) -> dict:
+        if episode_id not in self._meta:
+            path = self.window_file(episode_id, "meta.json", optional=True)
+            try:
+                self._meta[episode_id] = json.load(open(path)) if path else {}
+            except Exception:
+                self._meta[episode_id] = {}
+        return self._meta[episode_id]
+
+    # -- lifecycle ---------------------------------------------------------
+
+    def set_repo(self, repo: str, revision: str | None) -> None:
+        if (repo, revision) != (self.repo, self.revision):
+            self.reset()
+            self.repo, self.revision = repo, revision
+
+    def reset(self) -> None:
+        self._on_disk.clear()
+        self._pairs.clear()
+        self._composed.clear()
+        self._meta.clear()
+        self.rows = []
+        for d in (self.data, self.render):
+            shutil.rmtree(d, ignore_errors=True)
+            d.mkdir(parents=True, exist_ok=True)
+
+    def close(self) -> None:
+        self._pairs.clear()
+        self._composed.clear()
+        shutil.rmtree(self.root, ignore_errors=True)
+
+
+SESSIONS: dict[str, Session] = {}
+_LOCK = threading.Lock()
+
+
+def sweep() -> None:
+    """Drop sessions whose browser tab went away without a clean callback."""
+    now = time.time()
+    with _LOCK:
+        stale = [s for s in SESSIONS.values() if now - s.touched > SESSION_TTL]
+        for s in stale:
+            SESSIONS.pop(s.sid, None)
+            s.close()
+        live = set(SESSIONS)
+    for d in BASE_DIR.iterdir():  # dirs left behind by a previous process
+        try:
+            if d.is_dir() and d.name not in live and now - d.stat().st_mtime > SESSION_TTL:
+                shutil.rmtree(d, ignore_errors=True)
+        except OSError:
+            pass
+
+
+def get_session(sid: str | None) -> Session:
+    with _LOCK:
+        s = SESSIONS.get(sid or "")
+        if s is None:
+            sid = sid or uuid.uuid4().hex
+            s = SESSIONS[sid] = Session(sid)
+    s.touched = time.time()
+    return s
+
+
+def release(sid: str | None) -> None:
+    """gr.State delete_callback — fires when the session ends or its TTL expires."""
+    with _LOCK:
+        s = SESSIONS.pop(sid or "", None)
+    if s is not None:
+        s.close()
+
+
+# --------------------------------------------------------------------------- index
+
+
+def _row(order: int, episode_id: str, w: dict, m: dict) -> dict:
+    return {
+        "order": order,
+        "episode_id": episode_id,
+        "task": w.get("task") or "?",
+        "edge": w.get("edge_type") or "?",
+        "arm": w.get("arm") or "?",
+        "amplitude": w.get("amplitude"),
+        "start_frame": w.get("start_frame"),
+        "window": w.get("window"),
+        "critical_frame": w.get("critical_frame"),
+        "n_frames": w.get("n_frames"),
+        "full_psnr": m.get("full_psnr"),
+        "crit_psnr": m.get("crit_psnr"),
+        "full_mse": m.get("full_mse"),
+        "crit_mse": m.get("crit_mse"),
+    }
+
+
+def discover_episodes(session: Session) -> tuple[list[str], bool]:
+    """List windows/ one level deep, capped. Returns (ids, was_truncated)."""
+    tree = HfApi().list_repo_tree(
+        session.repo,
+        path_in_repo="windows",
+        recursive=False,
+        revision=session.revision,
         repo_type="dataset",
         token=HF_TOKEN,
     )
+    eps, truncated = [], False
+    for entry in tree:
+        if not isinstance(entry, RepoFolder):
+            continue
+        eps.append(entry.path.split("/")[-1])
+        if len(eps) >= MAX_DISCOVER:
+            truncated = True
+            break
+    return sorted(eps), truncated
+
+
+def build_index(session: Session) -> str:
+    """Populate session.rows from the manifest, or by listing the repo. Returns a note."""
+    metrics = []
+    path = session.fetch("metrics_summary.json", optional=True)
+    if path:
+        try:
+            metrics = json.load(open(path))
+        except Exception:
+            metrics = []
+    by_id = {m.get("id"): m for m in metrics if isinstance(m, dict)}
+
+    manifest = None
+    path = session.fetch("windows_manifest.json", optional=True)
+    if path:
+        try:
+            loaded = json.load(open(path))
+            if isinstance(loaded, list) and loaded:
+                manifest = loaded
+        except Exception:
+            manifest = None
+
+    if manifest is not None:
+        session.rows = [
+            _row(i, w.get("episode_id", f"window_{i}"), w, by_id.get(w.get("episode_id"), {}))
+            for i, w in enumerate(manifest)
+        ]
+        note = "`windows_manifest.json`"
+    else:
+        eps, truncated = discover_episodes(session)
+        if not eps:
+            raise FileNotFoundError(
+                "no windows/<episode_id>/ directories found — is this the right repo?"
+            )
+        session.rows = [_row(i, ep, {}, by_id.get(ep, {})) for i, ep in enumerate(eps)]
+        note = "directory listing (no `windows_manifest.json`)"
+        if truncated:
+            note += f", capped at {MAX_DISCOVER}"
 
+    if metrics:
+        note += " + `metrics_summary.json`"
+    return note
 
-@functools.lru_cache(maxsize=1)
-def load_index() -> list[dict]:
-    """Manifest rows joined with their metrics, in manifest order."""
-    manifest = json.load(open(fetch("windows_manifest.json")))
-    try:
-        metrics = json.load(open(fetch("metrics_summary.json")))
-    except Exception:
-        metrics = []
-    by_id = {m["id"]: m for m in metrics}
-
-    rows = []
-    for i, w in enumerate(manifest):
-        ep = w["episode_id"]
-        m = by_id.get(ep, {})
-        rows.append(
-            {
-                "order": i,
-                "episode_id": ep,
-                "task": w.get("task", "?"),
-                "edge": w.get("edge_type", "?"),
-                "arm": w.get("arm", "?"),
-                "amplitude": w.get("amplitude"),
-                "start_frame": w.get("start_frame"),
-                "window": w.get("window"),
-                "critical_frame": w.get("critical_frame"),
-                "n_frames": w.get("n_frames"),
-                "full_psnr": m.get("full_psnr"),
-                "crit_psnr": m.get("crit_psnr"),
-                "full_mse": m.get("full_mse"),
-                "crit_mse": m.get("crit_mse"),
-            }
-        )
-    return rows
+
+def short_id(episode_id: str) -> str:
+    return re.sub(r"^(episode|window|ep)[-_]", "", episode_id)[:8]
 
 
 def label_for(row: dict) -> str:
     psnr = row["crit_psnr"]
-    head = f"{psnr:5.1f} dB" if psnr is not None else "  -- dB"
-    return f"{head} · {row['edge']} {row['arm']} · {row['task']} · {row['episode_id'][8:16]}"
+    head = f"{psnr:5.1f} dB" if isinstance(psnr, (int, float)) else "  -- dB"
+    tags = " ".join(t for t in (row["edge"], row["arm"]) if t and t != "?")
+    parts = [head] + ([tags] if tags else []) + [row["task"], short_id(row["episode_id"])]
+    return " · ".join(parts)
 
 
-def filter_rows(query: str, edge: str, arm: str, sort: str) -> list[dict]:
-    rows = load_index()
+def filter_rows(session: Session, query: str, edge: str, arm: str, sort: str) -> list[dict]:
+    rows = session.rows
     q = (query or "").strip().lower()
     if q:
         rows = [r for r in rows if q in r["task"].lower() or q in r["episode_id"].lower()]
@@ -120,24 +360,35 @@ def filter_rows(query: str, edge: str, arm: str, sort: str) -> list[dict]:
     if arm != "any":
         rows = [r for r in rows if r["arm"] == arm]
 
-    worst = float("inf")
+    def psnr(r, missing):
+        v = r["crit_psnr"]
+        return v if isinstance(v, (int, float)) else missing
+
     if sort == SORTS[0]:
-        rows = sorted(rows, key=lambda r: r["crit_psnr"] if r["crit_psnr"] is not None else worst)
-    elif sort == SORTS[1]:
-        rows = sorted(rows, key=lambda r: -(r["crit_psnr"] if r["crit_psnr"] is not None else -worst))
-    elif sort == SORTS[2]:
-        rows = sorted(rows, key=lambda r: (r["task"], r["episode_id"]))
-    else:
-        rows = sorted(rows, key=lambda r: r["order"])
-    return rows
+        return sorted(rows, key=lambda r: psnr(r, float("inf")))
+    if sort == SORTS[1]:
+        return sorted(rows, key=lambda r: -psnr(r, float("-inf")))
+    if sort == SORTS[2]:
+        return sorted(rows, key=lambda r: (r["task"], r["episode_id"]))
+    return sorted(rows, key=lambda r: r["order"])
 
 
-@functools.lru_cache(maxsize=64)
-def row_by_id(episode_id: str) -> dict:
-    for r in load_index():
+def row_by_id(session: Session, episode_id: str) -> dict:
+    for r in session.rows:
         if r["episode_id"] == episode_id:
-            return r
-    raise KeyError(episode_id)
+            row = dict(r)
+            break
+    else:
+        row = _row(0, episode_id, {}, {})
+    # backfill whatever the manifest did not carry from the window's own meta.json
+    if row["task"] == "?" or row["critical_frame"] is None:
+        meta = session.meta(episode_id)
+        if meta:
+            merged = _row(row["order"], episode_id, meta, {})
+            for k, v in merged.items():
+                if row.get(k) in (None, "?"):
+                    row[k] = v
+    return row
 
 
 # --------------------------------------------------------------------------- video
@@ -158,18 +409,22 @@ def read_video(path: str) -> np.ndarray:
     return np.stack(frames)
 
 
-@functools.lru_cache(maxsize=4)
-def load_pair(episode_id: str) -> tuple[np.ndarray, np.ndarray]:
-    gt = read_video(fetch(f"windows/{episode_id}/gt.mp4"))
-    gen = read_video(fetch(f"windows/{episode_id}/generated.mp4"))
-    t = min(len(gt), len(gen))
-    return gt[:t], gen[:t]
+def n_cameras(frames: np.ndarray) -> int:
+    """How many roughly-square cameras are stacked vertically in the frame."""
+    h, w = frames.shape[1], frames.shape[2]
+    return max(1, min(6, int(round(h / w))))
+
 
+def camera_choices(n: int) -> list[tuple[str, int]]:
+    names = THREE_CAM_NAMES if n == 3 else [f"Camera {i + 1}" for i in range(n)]
+    choices = [(name, i) for i, name in enumerate(names)]
+    return ([("All cameras", ALL_CAMS)] + choices) if n > 1 else choices
 
-def cam_slices(frames: np.ndarray, view: str) -> list[np.ndarray]:
-    """Split the vstacked video into the camera bands requested by `view`."""
-    h = frames.shape[1] // 3
-    rows = [VIEW_ROW[view]] if view in VIEW_ROW else [0, 1, 2]
+
+def cam_slices(frames: np.ndarray, view: int) -> list[np.ndarray]:
+    n = n_cameras(frames)
+    h = frames.shape[1] // n
+    rows = range(n) if view == ALL_CAMS or not 0 <= view < n else [view]
     return [frames[:, i * h : (i + 1) * h] for i in rows]
 
 
@@ -185,7 +440,7 @@ def letterbox_rows(*bands: np.ndarray, thresh: int = 16) -> tuple[int, int]:
     return int(idx[0]), int(idx[-1]) + 1
 
 
-def crop_view(gt: np.ndarray, gen: np.ndarray, view: str) -> tuple[np.ndarray, np.ndarray]:
+def crop_view(gt: np.ndarray, gen: np.ndarray, view: int) -> tuple[np.ndarray, np.ndarray]:
     """Select cameras and trim the letterbox, identically for both videos."""
     gt_out, gen_out = [], []
     for g, p in zip(cam_slices(gt, view), cam_slices(gen, view)):
@@ -219,11 +474,9 @@ def diff_frames(gt: np.ndarray, gen: np.ndarray, gain: float = 3.0) -> np.ndarra
     return np.stack([cv2.applyColorMap(f, cv2.COLORMAP_INFERNO)[:, :, ::-1] for f in d])
 
 
-@functools.lru_cache(maxsize=2)
-def compose(episode_id: str, view: str, show_diff: bool, crit_idx: int) -> np.ndarray:
+def compose(pair: tuple[np.ndarray, np.ndarray], view: int, show_diff: bool, crit_idx: int) -> np.ndarray:
     """Build the labelled side-by-side strip: [T, H, W, 3] RGB."""
-    gt, gen = load_pair(episode_id)
-    gt_c, gen_c = crop_view(gt, gen, view)
+    gt_c, gen_c = crop_view(pair[0], pair[1], view)
     gt_v, gen_v = upscale(gt_c, SCALE), upscale(gen_c, SCALE)
 
     panels = [panel(gt_v, "GROUND TRUTH", GT_COLOR), panel(gen_v, "GENERATED", GEN_COLOR)]
@@ -240,14 +493,15 @@ def compose(episode_id: str, view: str, show_diff: bool, crit_idx: int) -> np.nd
     t, h, w, _ = comp.shape
     out = np.zeros((t, h + FTR, w, 3), np.uint8)
     out[:, :h] = comp
+    has_crit = 0 <= crit_idx < t
     tl0, tl1, y = int(w * 0.45), w - 14, FTR // 2
     for i in range(t):
         bar = np.full((FTR, w, 3), 22, np.uint8)
-        is_crit = i == crit_idx
+        is_crit = has_crit and i == crit_idx
         text = f"frame {i:02d}/{t - 1}" + ("   CRITICAL FRAME" if is_crit else "")
         cv2.putText(bar, text, (8, FTR - 8), FONT, 0.5, CRIT_COLOR if is_crit else (200, 200, 200), 1, cv2.LINE_AA)
         cv2.line(bar, (tl0, y), (tl1, y), (70, 70, 70), 2, cv2.LINE_AA)
-        if 0 <= crit_idx < t:
+        if has_crit:
             cx = int(tl0 + (tl1 - tl0) * crit_idx / max(t - 1, 1))
             cv2.line(bar, (cx, y - 7), (cx, y + 7), CRIT_COLOR, 2, cv2.LINE_AA)
         px = int(tl0 + (tl1 - tl0) * i / max(t - 1, 1))
@@ -264,7 +518,7 @@ def compose(episode_id: str, view: str, show_diff: bool, crit_idx: int) -> np.nd
 
 
 def write_mp4(frames: np.ndarray, path: Path, fps: float) -> str:
-    t, h, w, _ = frames.shape
+    _, h, w, _ = frames.shape
     writer = imageio_ffmpeg.write_frames(
         str(path),
         size=(w, h),
@@ -282,115 +536,178 @@ def write_mp4(frames: np.ndarray, path: Path, fps: float) -> str:
     return str(path)
 
 
-def render_video(episode_id: str, view: str, show_diff: bool, fps: float, crit_idx: int) -> str:
-    key = f"{episode_id}_{view.replace(' ', '')}_{int(show_diff)}_{fps:g}"
-    path = CACHE_DIR / f"{key}.mp4"
+def render_video(session: Session, episode_id: str, view: int, show_diff: bool, fps: float, crit: int) -> str:
+    path = session.render / f"{episode_id}_{view}_{int(show_diff)}_{fps:g}.mp4"
     if not path.exists():
-        write_mp4(compose(episode_id, view, show_diff, crit_idx), path, fps)
+        write_mp4(session.composed(episode_id, view, show_diff, crit), path, fps)
     return str(path)
 
 
 # --------------------------------------------------------------------------- actions plot
 
 
-def plot_actions(episode_id: str, crit_idx: int, cursor: int | None = None) -> Figure:
-    actions = np.asarray(json.load(open(fetch(f"windows/{episode_id}/actions.json"))), dtype=float)
+def plot_actions(session: Session, episode_id: str, crit_idx: int, cursor: int | None) -> Figure | None:
+    path = session.window_file(episode_id, "actions.json", optional=True)
+    if not path:
+        return None
+    try:
+        actions = np.atleast_2d(np.asarray(json.load(open(path)), dtype=float))
+    except Exception:
+        return None
+
     t = np.arange(len(actions))
+    dim = actions.shape[1]
+    if dim == 14:  # bimanual: [L arm 6, L grip, R arm 6, R grip]
+        specs = [
+            ("Grippers (1.0 = open, 0.0 = closed)", [(6, "L gripper"), (13, "R gripper")]),
+            ("Left arm joints", [(i, f"L{i}") for i in range(6)]),
+            ("Right arm joints", [(i, f"R{i - 7}") for i in range(7, 13)]),
+        ]
+    else:
+        specs = [(f"Action channels [{len(actions)} x {dim}]", [(i, f"a{i}") for i in range(dim)])]
 
-    fig = Figure(figsize=(7.2, 5.4), dpi=110, layout="constrained")
-    axes = fig.subplots(3, 1, sharex=True)
-    specs = [
-        ("Grippers (1.0 = open, 0.0 = closed)", [(6, "L gripper"), (13, "R gripper")]),
-        ("Left arm joints", [(i, f"L{i}") for i in range(6)]),
-        ("Right arm joints", [(i, f"R{i - 7}") for i in range(7, 13)]),
-    ]
+    fig = Figure(figsize=(7.2, 1.8 * len(specs) + 0.6), dpi=110, layout="constrained")
+    axes = np.atleast_1d(fig.subplots(len(specs), 1, sharex=True))
     for ax, (title, channels) in zip(axes, specs):
         for idx, name in channels:
             ax.plot(t, actions[:, idx], lw=1.4, label=name)
         ax.set_title(title, fontsize=9, loc="left")
-        ax.axvline(crit_idx, color="crimson", ls="--", lw=1.2, zorder=0)
+        if crit_idx >= 0:
+            ax.axvline(crit_idx, color="crimson", ls="--", lw=1.2, zorder=0)
         if cursor is not None:
             ax.axvline(cursor, color="0.35", ls=":", lw=1.0, zorder=0)
         ax.grid(alpha=0.25, lw=0.5)
         ax.tick_params(labelsize=8)
-        ax.legend(fontsize=7, ncol=6, loc="upper right", framealpha=0.7)
-    axes[-1].set_xlabel("frame in window (dashed red = critical frame)", fontsize=8)
+        if len(channels) <= 8:
+            ax.legend(fontsize=7, ncol=6, loc="upper right", framealpha=0.7)
+    tail = " (dashed red = critical frame)" if crit_idx >= 0 else ""
+    axes[-1].set_xlabel(f"frame in window{tail}", fontsize=8)
     return fig
 
 
 # --------------------------------------------------------------------------- ui glue
 
 
-def metrics_md(row: dict) -> str:
+def metrics_md(session: Session, row: dict) -> str:
     def fmt(v, unit=""):
         return f"{v:.2f}{unit}" if isinstance(v, (int, float)) else "—"
 
-    win = row.get("window") or [row["start_frame"], (row["start_frame"] or 0) + 32]
-
-    return (
-        f"### {row['task'].replace('_', ' ')}\n"
-        f"`{row['episode_id']}`\n\n"
-        f"| | full window | critical frame |\n|---|---|---|\n"
-        f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |\n"
-        f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |\n\n"
-        f"**{row['edge']}** · arm **{row['arm']}** · amplitude {fmt(row['amplitude'])} · "
-        f"window frames {win[0]}–{win[1]} of {row['n_frames']} · "
-        f"critical frame {row['critical_frame']}"
+    out = [f"### {row['task'].replace('_', ' ')}", f"`{row['episode_id']}`", ""]
+    if any(isinstance(row[k], (int, float)) for k in ("full_psnr", "crit_psnr", "full_mse", "crit_mse")):
+        out += [
+            "| | full window | critical frame |",
+            "|---|---|---|",
+            f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |",
+            f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |",
+            "",
+        ]
+    bits = []
+    if row["edge"] != "?":
+        bits.append(f"**{row['edge']}**")
+    if row["arm"] != "?":
+        bits.append(f"arm **{row['arm']}**")
+    if isinstance(row["amplitude"], (int, float)):
+        bits.append(f"amplitude {fmt(row['amplitude'])}")
+    win = row.get("window")
+    if isinstance(win, (list, tuple)) and len(win) == 2:
+        bits.append(f"window frames {win[0]}–{win[1]}")
+    if row["n_frames"]:
+        bits.append(f"of {row['n_frames']}")
+    if row["critical_frame"] is not None:
+        bits.append(f"critical frame {row['critical_frame']}")
+    if bits:
+        out.append(" · ".join(bits))
+    out.append(
+        f"\n<sub>cached this session: {session.disk_windows()}/{MAX_WINDOWS_ON_DISK} windows "
+        "— evicted least-recently-used, all deleted when you close the tab</sub>"
     )
+    return "\n".join(out)
 
 
 def table_rows(rows: list[dict]) -> list[list]:
+    def num(v):
+        return round(v, 2) if isinstance(v, (int, float)) else None
+
     return [
-        [
-            r["episode_id"][8:16],
-            r["task"],
-            r["edge"],
-            r["arm"],
-            round(r["crit_psnr"], 2) if r["crit_psnr"] is not None else None,
-            round(r["full_psnr"], 2) if r["full_psnr"] is not None else None,
-        ]
+        [short_id(r["episode_id"]), r["task"], r["edge"], r["arm"], num(r["crit_psnr"]), num(r["full_psnr"])]
         for r in rows
     ]
 
 
-def on_filter(query, edge, arm, sort, current):
-    rows = filter_rows(query, edge, arm, sort)
+def critical_index(row: dict) -> int:
+    """Offset of the critical frame inside the window, or NO_CRIT."""
+    start, crit = row.get("start_frame"), row.get("critical_frame")
+    if start is None or crit is None:
+        return NO_CRIT
+    return int(crit) - int(start)
+
+
+def on_load(repo_text, sid):
+    """Resolve a repo and load only its index — no window data is pulled here."""
+    sweep()
+    session = get_session(sid)
+    repo, revision = parse_repo(repo_text)
+    if not repo:
+        return session.sid, "Enter a dataset repo id, e.g. `owner/name`."
+    session.set_repo(repo, revision)
+    try:
+        note = build_index(session)
+    except Exception as exc:
+        session.rows = []
+        gr.Warning(f"Could not load {repo}: {exc}")
+        return session.sid, f"❌ **{repo}** — {exc}"
+    at = f"@{revision}" if revision else ""
+    return session.sid, (
+        f"✅ **{repo}{at}** — {len(session.rows)} windows from {note}. "
+        "Window files download only when you open them."
+    )
+
+
+def on_filter(sid, query, edge, arm, sort, current):
+    session = get_session(sid)
+    if not session.rows:
+        return gr.update(choices=[], value=None), [], []
+    rows = filter_rows(session, query, edge, arm, sort)
     choices = [(label_for(r), r["episode_id"]) for r in rows]
     ids = [r["episode_id"] for r in rows]
     value = current if current in ids else (ids[0] if ids else None)
     return (
-        gr.update(choices=choices, value=value, label=f"Window ({len(ids)} of 300)"),
+        gr.update(choices=choices, value=value, label=f"Window ({len(ids)} of {len(session.rows)})"),
         table_rows(rows),
         ids,
     )
 
 
-def critical_index(row: dict) -> int:
-    """Offset of the critical frame inside the window (16 for 33-frame windows)."""
-    start, crit = row.get("start_frame"), row.get("critical_frame")
-    if start is None or crit is None:
-        return 16
-    return int(crit) - int(start)
+def on_select(sid, episode_id, view, show_diff, fps, frame_idx):
+    session = get_session(sid)
+    if not session.rows or not episode_id:
+        return None, None, "*No window selected.*", None, gr.update(), gr.update()
+    row = row_by_id(session, episode_id)
+    crit = critical_index(row)
 
+    cams = n_cameras(session.pair(episode_id)[0])
+    choices = camera_choices(cams)
+    valid = [v for _, v in choices]
+    view = view if view in valid else choices[0][1]
+    cam_update = gr.update(choices=choices, value=view)
 
-def on_select(episode_id, view, show_diff, fps, frame_idx):
-    if not episode_id:
-        return None, None, "*No window matches the current filters.*", None
-    row = row_by_id(episode_id)
-    crit = critical_index(row)
-    video = render_video(episode_id, view, show_diff, fps, crit)
-    frames = compose(episode_id, view, show_diff, crit)
-    i = int(min(frame_idx, len(frames) - 1))
-    return video, frames[i], metrics_md(row), plot_actions(episode_id, crit, i)
+    video = render_video(session, episode_id, view, show_diff, fps, crit)
+    frames = session.composed(episode_id, view, show_diff, crit)
+    last = len(frames) - 1
+    i = int(min(frame_idx, last))
+    plot = plot_actions(session, episode_id, crit, i)
+    slider = gr.update(maximum=last, value=i, label=f"Frame in window (0–{last})")
+    return video, frames[i], metrics_md(session, row), plot, slider, cam_update
 
 
-def on_frame(episode_id, view, show_diff, frame_idx):
-    if not episode_id:
+def on_frame(sid, episode_id, view, show_diff, frame_idx):
+    session = get_session(sid)
+    if not session.rows or not episode_id:
         return None, None
-    crit = critical_index(row_by_id(episode_id))
-    frames = compose(episode_id, view, show_diff, crit)
+    crit = critical_index(row_by_id(session, episode_id))
+    frames = session.composed(episode_id, view, show_diff, crit)
     i = int(min(frame_idx, len(frames) - 1))
-    return frames[i], plot_actions(episode_id, crit, i)
+    return frames[i], plot_actions(session, episode_id, crit, i)
 
 
 def step(ids, current, delta):
@@ -410,15 +727,23 @@ CSS = """
 .window-still img { max-height: 78vh; object-fit: contain; background: #111; }
 """
 
-with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
-    gr.Markdown(
-        "# Forward-dynamics rollout vs. ground truth\n"
-        f"Ground-truth and generated 33-frame windows from **`{DATASET_REPO}`**, "
-        "each centred on a gripper critical frame (frame 16). "
-        "Cameras are top / left-wrist / right-wrist."
-    )
+with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(600, 600)) as demo:
+    gr.Markdown("# Forward-dynamics rollout vs. ground truth")
+
+    sid_state = gr.State("", time_to_live=SESSION_TTL, delete_callback=release)
     ids_state = gr.State([])
 
+    with gr.Row():
+        repo_box = gr.Textbox(
+            DEFAULT_REPO,
+            label="Dataset repo",
+            info="owner/name, optionally owner/name@revision. A repo URL also works.",
+            scale=5,
+            submit_btn=True,
+        )
+        load_btn = gr.Button("Load", variant="primary", scale=1)
+    status = gr.Markdown()
+
     with gr.Row():
         with gr.Column(scale=1, min_width=320):
             query = gr.Textbox(label="Search task or episode", placeholder="e.g. napkin, pillowcase…")
@@ -431,9 +756,9 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
                 prev_btn = gr.Button("← Prev")
                 next_btn = gr.Button("Next →")
             gr.Markdown("---")
-            view = gr.Radio(VIEWS, value="Top", label="Camera")
+            view = gr.Radio(camera_choices(3), value=0, label="Camera")
             show_diff = gr.Checkbox(False, label="Add |GT − generated| difference panel")
-            fps = gr.Slider(2, 30, value=8, step=1, label="Playback fps (source is 30 fps)")
+            fps = gr.Slider(2, 30, value=8, step=1, label="Playback fps")
             info = gr.Markdown()
 
         with gr.Column(scale=2):
@@ -447,10 +772,10 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
                         elem_classes="window-video",
                     )
                 with gr.Tab("Frame stepper"):
-                    frame_idx = gr.Slider(0, 32, value=16, step=1, label="Frame in window (16 = critical)")
+                    frame_idx = gr.Slider(0, 32, value=16, step=1, label="Frame in window")
                     still = gr.Image(label="Frame comparison", elem_classes="window-still")
                 with gr.Tab("Driving actions"):
-                    actions_plot = gr.Plot(label="Action chunk fed to the model [33 x 14]")
+                    actions_plot = gr.Plot(label="Driving action chunk")
                 with gr.Tab("All windows"):
                     table = gr.Dataframe(
                         headers=["episode", "task", "edge", "arm", "crit PSNR", "full PSNR"],
@@ -460,12 +785,23 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
                         max_height=650,
                     )
 
-    filter_inputs = [query, edge, arm, sort]
+    filter_inputs = [sid_state, query, edge, arm, sort]
     filter_outputs = [picker, table, ids_state]
-    view_inputs = [picker, view, show_diff, fps, frame_idx]
-    view_outputs = [video, still, info, actions_plot]
+    view_inputs = [sid_state, picker, view, show_diff, fps, frame_idx]
+    view_outputs = [video, still, info, actions_plot, frame_idx, view]
+
+    def wire_load(event):
+        return (
+            event(on_load, [repo_box, sid_state], [sid_state, status])
+            .then(on_filter, filter_inputs + [picker], filter_outputs)
+            .then(on_select, view_inputs, view_outputs)
+        )
 
-    for control in filter_inputs:
+    wire_load(load_btn.click)
+    wire_load(repo_box.submit)
+    wire_load(demo.load)
+
+    for control in (query, edge, arm, sort):
         control.change(on_filter, filter_inputs + [picker], filter_outputs).then(
             on_select, view_inputs, view_outputs
         )
@@ -473,7 +809,7 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
     picker.change(on_select, view_inputs, view_outputs)
     for control in (view, show_diff, fps):
         control.change(on_select, view_inputs, view_outputs)
-    frame_idx.change(on_frame, [picker, view, show_diff, frame_idx], [still, actions_plot])
+    frame_idx.change(on_frame, [sid_state, picker, view, show_diff, frame_idx], [still, actions_plot])
 
     prev_btn.click(lambda i, c: step(i, c, -1), [ids_state, picker], picker)
     next_btn.click(lambda i, c: step(i, c, +1), [ids_state, picker], picker)
@@ -483,9 +819,5 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True) as demo:
 
     table.select(on_table_select, ids_state, picker)
 
-    demo.load(on_filter, filter_inputs + [picker], filter_outputs).then(
-        on_select, view_inputs, view_outputs
-    )
-
 if __name__ == "__main__":
     demo.queue(max_size=16).launch(css=CSS)