add wandb + aiohttp to requirements; drop local dev tooling

assemble.sh

@@ -1,85 +0,0 @@
-#!/usr/bin/env bash
-# Stage everything the Space needs into space/, then push it to Hugging Face.
-#
-# The Space is its own git repo, so this copies the pieces it needs out of the
-# main repo rather than depending on it. Re-run after changing anything under
-# src/ and push again.
-#
-# Usage:
-#   ./assemble.sh                      # stage only
-#   ./assemble.sh <user-or-org>/<space-name>    # stage, then show the push commands
-#
-# Env overrides:
-#   IDM_SRC   IDM checkpoint to ship (default: base_run3/step_040000.pt)
-set -euo pipefail
-cd "$(dirname "$0")"
-
-REPO_ROOT="$(cd .. && pwd)"
-IDM_SRC="${IDM_SRC:-$REPO_ROOT/src/logs/idm_tokenizer_ckpts/base_run3/step_040000.pt}"
-SPACE_ID="${1:-}"
-
-echo "staging from $REPO_ROOT"
-
-# --- code -------------------------------------------------------------------
-rm -rf src
-mkdir -p src
-# Flat modules the server imports; envs/ is a package and comes wholesale.
-cp "$REPO_ROOT"/src/*.py src/
-cp -r "$REPO_ROOT"/src/envs src/
-# Only the v2 page is served, but the server's loader falls back by filename, so
-# ship the one it is pointed at and nothing else.
-cp "$REPO_ROOT"/src/interactive_corr_v2.html src/
-# Language embeddings, read at runtime for the task conditioning.
-cp "$REPO_ROOT"/tasks.json .
-
-# Training-only artefacts must not ride along.
-rm -rf src/logs src/checkpoints src/wandb src/data
-find src -name '__pycache__' -type d -prune -exec rm -rf {} +
-
-# --- weights ----------------------------------------------------------------
-mkdir -p assets
-if [ ! -f "$IDM_SRC" ]; then
-    echo "ERROR: IDM checkpoint not found: $IDM_SRC" >&2
-    echo "       set IDM_SRC=<path> or train one with src/run_idm_tokenizer.sh" >&2
-    exit 1
-fi
-python strip_idm.py "$IDM_SRC" assets/idm_tokenizer.pt
-
-# The tokenizer/dynamics pair is public on the Hub and downloaded at startup, so
-# it is deliberately NOT vendored here — it would add 1.4 GB to every push.
-
-# --- git config -------------------------------------------------------------
-cat > .gitattributes <<'EOF'
-*.pt filter=lfs diff=lfs merge=lfs -text
-tasks.json filter=lfs diff=lfs merge=lfs -text
-EOF
-
-cat > .gitignore <<'EOF'
-checkpoints/
-__pycache__/
-*.pyc
-EOF
-
-echo
-echo "staged:"
-printf '  %-28s %s\n' \
-    "app.py"                 "$(du -h app.py | cut -f1)" \
-    "src/ (+envs)"           "$(du -sh src | cut -f1)" \
-    "assets/idm_tokenizer.pt" "$(du -h assets/idm_tokenizer.pt | cut -f1)" \
-    "tasks.json"             "$(du -h tasks.json | cut -f1)"
-
-if [ -n "$SPACE_ID" ]; then
-    cat <<EOF
-
-to publish:
-  git init && git lfs install
-  git remote add origin https://huggingface.co/spaces/$SPACE_ID
-  git add -A && git commit -m "deploy correlation demo"
-  git push -u origin main
-
-the Space needs hardware = ZeroGPU (Settings -> Hardware).
-EOF
-else
-    echo
-    echo "pass a Space id (e.g. ./assemble.sh myname/wm-hallucination) for push commands."
-fi

bench.py

@@ -1,89 +0,0 @@
-"""Drive a running demo over its WebSocket and report where the time goes.
-
-The server already timestamps each render step and ships it as `ms` in the status
-frame, so this measures the real per-step cost rather than round-trip latency.
-Steps that also ran the WAV rollout are reported separately: that is the one
-periodic cost, and averaging it into every step hides which knob to turn.
-
-Usage:
-    python bench.py                       # 60 steps against localhost:7860
-    python bench.py --steps 120 --port 7860
-    python bench.py --task walker-run     # switch task first
-"""
-import argparse
-import asyncio
-import json
-import statistics as st
-
-import aiohttp
-
-
-async def run(args) -> int:
-    url = f"http://{args.host}:{args.port}"
-    async with aiohttp.ClientSession() as s:
-        async with s.get(f"{url}/healthz", timeout=aiohttp.ClientTimeout(total=10)) as r:
-            print((await r.text()).strip())
-
-        async with s.ws_connect(f"{url}/ws") as ws:
-            if args.task:
-                await ws.send_json({"type": "set_task", "task": args.task})
-                await asyncio.sleep(2.0)
-            await ws.send_json({"type": "keydown", "key": "ArrowRight"})
-
-            plain, with_wav, n_jpeg, first, last = [], [], 0, None, None
-            while len(plain) + len(with_wav) < args.steps:
-                try:
-                    m = await asyncio.wait_for(ws.receive(), timeout=60)
-                except asyncio.TimeoutError:
-                    print("timed out waiting for frames")
-                    break
-                if m.type == aiohttp.WSMsgType.BINARY:
-                    n_jpeg += 1
-                    continue
-                if m.type != aiohttp.WSMsgType.TEXT:
-                    break
-                d = json.loads(m.data)
-                if d.get("type") != "status" or "ms" not in d:
-                    continue
-                if first is None:
-                    first = d.get("step", 0)
-                last = d.get("step", 0)
-                (with_wav if d.get("wav_steps") else plain).append(float(d["ms"]))
-            await ws.send_json({"type": "disconnect"})
-
-    every = len(plain) + len(with_wav)
-    if not every:
-        print("no timed frames received")
-        return 1
-
-    def line(name, xs):
-        if not xs:
-            print(f"  {name:<22s}  (none)")
-            return
-        xs = sorted(xs)
-        print(f"  {name:<22s} n={len(xs):>4d}  median={st.median(xs):7.1f} ms   "
-              f"p90={xs[int(0.9 * (len(xs) - 1))]:7.1f}   max={xs[-1]:7.1f}")
-
-    print(f"\nper-step render cost   (steps {first} -> {last}, {n_jpeg} jpeg frames)")
-    line("plain step", plain)
-    line("step + WAV rollout", with_wav)
-
-    allms = plain + with_wav
-    mean_ms = st.mean(allms)
-    print(f"\n  mean over all steps    {mean_ms:7.1f} ms  ->  {1000 / mean_ms:5.2f} fps sustained")
-    if plain and with_wav:
-        extra = st.median(with_wav) - st.median(plain)
-        share = 100.0 * extra / st.median(with_wav)
-        print(f"  WAV rollout adds       {extra:7.1f} ms on the steps it runs ({share:.0f}% of them)")
-        print(f"  amortised over 1-in-{round(len(allms) / max(1, len(with_wav)))} steps: "
-              f"{extra * len(with_wav) / len(allms):.1f} ms/step")
-    return 0
-
-
-if __name__ == "__main__":
-    p = argparse.ArgumentParser()
-    p.add_argument("--host", default="127.0.0.1")
-    p.add_argument("--port", type=int, default=7860)
-    p.add_argument("--steps", type=int, default=60)
-    p.add_argument("--task", type=str, default=None)
-    raise SystemExit(asyncio.run(run(p.parse_args())))

requirements.txt

@@ -15,8 +15,17 @@ gradio==6.17.3
 spaces
 huggingface_hub
 websockets
+# The interactive_* modules import `aiohttp.web` at module scope (they ship their
+# own aiohttp server; the Space only reuses their classes). It reaches the image
+# solely through extras nobody selects -- huggingface_hub[inference], fsspec[http] --
+# so it has to be requested outright.
+aiohttp
 
 # world model / data
+# wandb is a *training* dependency, but train_dynamics.py and idm_tokenizer.py
+# import it at module scope, and the demo pulls sampler helpers out of both. It is
+# declared only under ogbench's "train" extra, hence not installed by default.
+wandb
 transformers==4.56.2
 lpips==0.1.4
 tensordict==0.10.0

run_local.sh

@@ -1,78 +0,0 @@
-#!/usr/bin/env bash
-# Run the Space app locally, exactly as it will run on HF but with a real GPU.
-#
-#   ./run_local.sh              # balanced (the Space default)
-#   ./run_local.sh fast         # ~4x faster, for interaction; slightly coarser sampling
-#   ./run_local.sh faithful     # matches the offline experiments; slow
-#
-# Then, in another shell:
-#   python bench.py             # where the per-step time actually goes
-#
-# Env overrides: GPU, PORT, CKPT_DIR, plus any knob below.
-set -euo pipefail
-cd "$(dirname "$0")"
-
-PRESET="${1:-balanced}"
-GPU="${GPU:-0}"
-PORT="${PORT:-7860}"
-# Reuse an existing checkpoint tree if there is one, rather than re-downloading.
-CKPT_DIR="${CKPT_DIR:-$(cd .. && pwd)/src/checkpoints}"
-
-case "$PRESET" in
-  fast)
-    # The two dominant costs are the Euler step count (every sampler call pays
-    # it) and the WAV rollout (H steps x 2 action sources). KV caching is exact
-    # -- same numbers, less compute -- so it is on in every preset.
-    EVAL_D=0.5            # 2 Euler steps per sample instead of 8
-    CTX_WINDOW=12         # attention context; must stay >= wav_context + wav_horizon_long - 1
-    WAV_HORIZON_LONG=4
-    WAV_EVERY=4
-    N_SAMPLES_U=2
-    U_EVERY=2             # tokenizer round-trip (u_r) every other frame
-    ;;
-  balanced)
-    EVAL_D=0.25           # 4 Euler steps
-    CTX_WINDOW=16
-    WAV_HORIZON_LONG=8
-    WAV_EVERY=4
-    N_SAMPLES_U=2
-    U_EVERY=1
-    ;;
-  faithful)
-    EVAL_D=0.125          # 8 Euler steps — the repo default
-    CTX_WINDOW=24
-    WAV_HORIZON_LONG=8
-    WAV_EVERY=4
-    N_SAMPLES_U=4
-    U_EVERY=1
-    ;;
-  *)
-    echo "unknown preset '$PRESET' (expected: fast | balanced | faithful)" >&2
-    exit 2
-    ;;
-esac
-
-# ctx_window bounds how much latent history a session keeps, and the lagged
-# rollout needs wav_context + wav_horizon_long of it. Fail loudly here rather
-# than letting the correlation rows sit empty forever.
-NEED=$(( 8 + WAV_HORIZON_LONG ))
-if [ "$(( CTX_WINDOW + 1 ))" -lt "$NEED" ]; then
-    echo "preset '$PRESET' is inconsistent: ctx_window=$CTX_WINDOW keeps $(( CTX_WINDOW + 1 ))" >&2
-    echo "latents but the rollout needs $NEED. Raise CTX_WINDOW or lower WAV_HORIZON_LONG." >&2
-    exit 1
-fi
-
-echo "preset=$PRESET  gpu=$GPU  port=$PORT"
-echo "  eval_d=$EVAL_D (Euler steps=$(python -c "print(int(1/$EVAL_D))"))  ctx_window=$CTX_WINDOW"
-echo "  wav: horizon_long=$WAV_HORIZON_LONG every=$WAV_EVERY   u: n_samples=$N_SAMPLES_U every=$U_EVERY"
-echo "  kv_cache=on"
-echo
-
-exec env \
-    CUDA_VISIBLE_DEVICES="$GPU" \
-    CKPT_DIR="$CKPT_DIR" \
-    WAV_EVERY="$WAV_EVERY" \
-    WAV_HORIZON_LONG="$WAV_HORIZON_LONG" \
-    EXTRA_ARGS="--eval_d $EVAL_D --ctx_window $CTX_WINDOW --n_samples_u $N_SAMPLES_U --u_every $U_EVERY --kv_cache" \
-    PORT="$PORT" \
-    python app.py

strip_idm.py

@@ -1,46 +0,0 @@
-"""Strip optimizer state from an IDM checkpoint before shipping it.
-
-`idm_tokenizer.py` saves AdamW's moments alongside the weights so a `--resume`
-is exact. That triples the file (252 MB vs 84 MB for the d_hidden=1024 head) and
-is dead weight in a deployment, which only ever loads the head for inference.
-
-`load_idm_tokenizer_head` recovers the architecture from the tensor shapes, so
-dropping everything but `model` (plus `args`/`step` for provenance) is safe.
-
-Usage:
-    python strip_idm.py <src.pt> <dst.pt>
-"""
-import os
-import sys
-
-import torch
-
-
-def main() -> int:
-    if len(sys.argv) != 3:
-        print(__doc__.strip(), file=sys.stderr)
-        return 2
-    src, dst = sys.argv[1], sys.argv[2]
-    if not os.path.isfile(src):
-        print(f"ERROR: no such checkpoint: {src}", file=sys.stderr)
-        return 1
-
-    ck = torch.load(src, map_location="cpu", weights_only=False)
-    if "model" not in ck:
-        print(f"ERROR: {src} has no 'model' key (got {sorted(ck)})", file=sys.stderr)
-        return 1
-
-    lean = {"model": ck["model"], "args": ck.get("args"), "step": ck.get("step")}
-    os.makedirs(os.path.dirname(os.path.abspath(dst)) or ".", exist_ok=True)
-    torch.save(lean, dst)
-
-    a, b = os.path.getsize(src) / 1e6, os.path.getsize(dst) / 1e6
-    dropped = sorted(k for k in ck if k not in lean)
-    print(f"{src}  {a:.0f} MB")
-    print(f"{dst}  {b:.0f} MB   (-{100 * (1 - b / a):.0f}%, dropped: {dropped or 'nothing'})")
-    print(f"step={lean['step']}  params={sum(v.numel() for v in lean['model'].values()) / 1e6:.1f}M")
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

test_page.py

@@ -1,214 +0,0 @@
-"""Execute the served page's JavaScript against captured server frames.
-
-Static checks (does it parse? are all identifiers declared?) kept missing real
-bugs here, because the page's failure mode is an exception thrown inside the
-WebSocket message handler — which the browser swallows, leaving panels frozen
-while data keeps arriving. The only way to catch that is to *run* the script.
-
-So: a minimal DOM in QuickJS, the real page script loaded into it, and real
-status frames replayed through the captured `ws.onmessage`. Then assert the
-panels actually changed.
-
-Usage:
-    python test_page.py [--url http://127.0.0.1:7860] [--frames /tmp/frames.json]
-    python test_page.py --html /path/to/page.html --frames /tmp/frames.json
-"""
-import argparse
-import json
-import re
-import sys
-import urllib.request
-
-import quickjs
-
-DOM = r"""
-var __errors = [];
-var __els = {};       // id -> element
-var __created = [];   // every element made by script
-// Ids that actually exist in the served markup. getElementById must return null
-// for anything else -- auto-creating a stand-in is precisely how a stub hides
-// the bug it is supposed to catch (a real browser returns null and the next
-// property write throws).
-var __realIds = {};
-
-function __mkEl(tag, id) {
-  var e = {
-    tagName: tag, id: id || "", className: "", textContent: "", innerHTML: "",
-    style: {}, dataset: {}, children: [], value: "", max: "", min: "",
-    appendChild: function (c) { this.children.push(c); return c; },
-    setAttribute: function (k, v) { this[k] = v; },
-    getAttribute: function (k) { return this[k]; },
-    addEventListener: function () {},
-    removeEventListener: function () {},
-    getBoundingClientRect: function () { return {width: 300, height: 120, top: 0, left: 0, right: 300, bottom: 120}; },
-    getContext: function () { return {}; },
-    querySelectorAll: function () { return []; },
-    classList: {add: function(){}, remove: function(){}, contains: function(){return false;}},
-    focus: function () {}, blur: function () {},
-  };
-  __created.push(e);
-  return e;
-}
-
-var document = {
-  documentElement: {
-    getAttribute: function () { return null; },
-    setAttribute: function () {},
-    hasAttribute: function () { return false; },
-    style: {},
-  },
-  getElementById: function (id) {
-    if (!__realIds[id]) return null;          // faithful to the browser
-    if (!__els[id]) __els[id] = __mkEl("div", id);
-    return __els[id];
-  },
-  createElement: function (t) { return __mkEl(t); },
-  createElementNS: function (ns, t) { return __mkEl(t); },
-  createTextNode: function (s) { return {textContent: String(s), nodeValue: String(s)}; },
-  createDocumentFragment: function () { return __mkEl("fragment"); },
-  // Selector support is deliberately shallow: the page only uses it to relabel
-  // the table header. Return a stable array per selector so index access works
-  // the way it does in a browser instead of blowing up on undefined.
-  __qsa: {},
-  querySelectorAll: function (sel) {
-    if (!this.__qsa[sel]) {
-      var n = 8, a = [];
-      for (var i = 0; i < n; i++) a.push(__mkEl("th"));
-      this.__qsa[sel] = a;
-    }
-    return this.__qsa[sel];
-  },
-  addEventListener: function () {},
-};
-
-var window = {
-  addEventListener: function () {},
-  innerWidth: 1200, innerHeight: 800, devicePixelRatio: 1,
-};
-var location = {protocol: "http:", host: "127.0.0.1:7860"};
-function getComputedStyle() { return {getPropertyValue: function () { return "#3987e5"; }}; }
-function matchMedia() { return {matches: false}; }
-function setTimeout(f) { return 0; }
-function clearTimeout() {}
-var console = {log: function(){}, error: function(){ __errors.push(Array.prototype.join.call(arguments, " ")); },
-               warn: function(){}};
-var URL = {createObjectURL: function () { return "blob:x"; }, revokeObjectURL: function () {}};
-function Blob() {}
-
-// Capture the socket the page opens so frames can be pushed through it.
-var __ws = null;
-function WebSocket(u) { __ws = this; this.readyState = 1; this.url = u; }
-WebSocket.OPEN = 1;
-WebSocket.prototype.send = function () {};
-WebSocket.prototype.close = function () {};
-
-function __feed(json) {
-  if (!__ws || !__ws.onmessage) { __errors.push("no onmessage handler"); return; }
-  __ws.onmessage({data: json});
-}
-function __panelValues() {
-  var out = [];
-  for (var i = 0; i < __created.length; i++) {
-    var e = __created[i];
-    if (e.className === "pCum") {
-      out.push(e.className + ":" + e.textContent);
-    }
-  }
-  return JSON.stringify(out);
-}
-// Exercise a slider the way a user does: set its value, fire its handler.
-// A control that renders but is not wired reads as "the numbers never change".
-function __slide(id, value) {
-  var el = __els[id];
-  if (!el) { __errors.push("no such control: " + id); return false; }
-  if (!el.oninput) { __errors.push("control has no oninput: " + id); return false; }
-  el.value = String(value);
-  el.oninput();
-  return true;
-}
-function __errCount() { return __errors.length; }
-function __errText() { return JSON.stringify(__errors.slice(0, 5)); }
-"""
-
-
-def main() -> int:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--url", default="http://127.0.0.1:7860")
-    ap.add_argument("--html", default=None)
-    ap.add_argument("--frames", default="/tmp/frames.json")
-    a = ap.parse_args()
-
-    html = (open(a.html).read() if a.html
-            else urllib.request.urlopen(a.url, timeout=10).read().decode())
-    js = re.search(r"<script>(.*)</script>", html, re.S).group(1)
-    frames = json.load(open(a.frames))
-
-    ctx = quickjs.Context()
-    ctx.eval(DOM)
-    # Teach the stub which ids the markup really has.
-    real_ids = sorted(set(re.findall(r'id="([A-Za-z_0-9]+)"', html)))
-    ctx.eval("__realIds = " + json.dumps({i: True for i in real_ids}) + ";")
-    print(f"markup declares {len(real_ids)} element ids")
-    try:
-        ctx.eval(js)
-    except Exception as e:
-        print(f"FAIL: the page script threw at load:\n  {e}")
-        return 1
-    print("page script loaded without throwing")
-
-    before = ctx.eval("__panelValues()")
-    n_samples = 0
-    for f in frames:
-        ctx.eval(f"__feed({json.dumps(json.dumps(f))})")
-        if f.get("wav_steps"):
-            n_samples += 1
-    after = ctx.eval("__panelValues()")
-
-    n_err = ctx.eval("__errCount()")
-    print(f"replayed {len(frames)} frames ({n_samples} carrying rollout samples)")
-
-    ok = True
-    if n_err:
-        print(f"FAIL: {n_err} error(s) surfaced from the message handler:")
-        for e in json.loads(ctx.eval("__errText()")):
-            print(f"    {e}")
-        ok = False
-
-    # Controls must actually rescore what is already on screen.
-    for cid, val in (("nsRange", 20), ("hRange", 6)):
-        moved_from = ctx.eval("__panelValues()")
-        if not ctx.eval(f'__slide("{cid}", {val})'):
-            print(f"FAIL: control {cid} is not wired")
-            ok_ctl = False
-            continue
-        moved_to = ctx.eval("__panelValues()")
-        if moved_from == moved_to:
-            print(f"WARN: moving {cid} to {val} changed nothing "
-                  f"(may be legitimate if the data is degenerate)")
-        else:
-            print(f"control {cid} -> {val}: panels rescored")
-
-    b, af = json.loads(before), json.loads(after)
-    changed = [x for x in af if x not in b]
-    if not af:
-        print("FAIL: no panel elements were created — buildPanels() never ran")
-        ok = False
-    elif not changed:
-        print("FAIL: panels never changed after replaying samples")
-        print(f"    still: {af[:8]}")
-        ok = False
-    else:
-        filled = [x for x in af if not x.endswith(":—") and not x.endswith(":")]
-        print(f"panels updated: {len(changed)} of {len(af)} slots changed, "
-              f"{len(filled)} now hold a value")
-        print(f"    sample: {filled[:6]}")
-        if not filled:
-            print("FAIL: every panel slot is still empty")
-            ok = False
-
-    print("\nPASS" if ok else "\nFAIL")
-    return 0 if ok else 1
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())