Cold-start progress, early WebGPU bail, and Hugging Face Space packaging

index.html

@@ -38,7 +38,7 @@
         var q = window.location.search || "";
         var go =
           /^#?r=/.test(h) ||
-          /[?&]mode=science\b/.test(q) ||
+          /[?&]mode=(science|game)\b/.test(q) ||
           /[?&](play|app)\b/.test(q);
         if (go) window.location.replace("app.html" + q + h);
       })();

package.json

@@ -23,6 +23,7 @@
     "bench:brain": "playwright test tests/bench.spec.ts --reporter=list",
     "build:slim": "npm run build && rm -rf dist/flybody dist/flybody.bundle.bin dist/brain.bin dist/brain.meta.json dist/vnc.bin dist/vnc.meta.json dist/walking-policy.bin dist/walking-obs-norm.bin dist/walking-ref.bin dist/walking-policy-fixtures.json",
     "deploy": "npm run build:slim && npx --yes wrangler pages deploy dist --project-name=webgpu-fly --branch=main",
+    "deploy:hf": "npm run build:slim && cp space/README.md space/.gitattributes dist/ && rm -f dist/_headers && hf upload abgunaydin/webgpu-fly dist . --repo-type=space",
     "deploy:vercel": "npm run build && vercel --prod",
     "deploy:r2": "bash tools/upload_to_r2.sh"
   },

space/.gitattributes

@@ -0,0 +1,45 @@
+# Hugging Face Space .gitattributes.
+#
+# This is the Hub default MINUS *.wasm and *.bin. Those two lines matter here:
+# LFS/Xet-backed files are served as a 302 to a per-request signed URL with
+# `cache-control: no-store`, so the browser cache never hits and dist's 8.6 MB
+# mujoco wasm would re-download on every single page load. Served as a plain
+# file it comes back inline with an ETag and revalidates as 304. It is under
+# the forced-LFS threshold, so this is allowed.
+#
+# The large connectome binaries are NOT in this repo at all — they live on R2
+# and are fetched at runtime, so nothing here needs LFS for them either.
+
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

space/README.md

@@ -0,0 +1,64 @@
+---
+title: webgpu-fly
+emoji: 🪰
+colorFrom: indigo
+colorTo: purple
+sdk: static
+app_file: index.html
+header: mini
+pinned: false
+license: mit
+short_description: A fly connectome and a physics fly body, live in WebGPU
+---
+
+# webgpu-fly
+
+**139,255 FlyWire brain neurons and 23,188 Janelia MANC ventral-nerve-cord
+neurons, simulated as leaky integrate-and-fire networks on your GPU, next to a
+TuragaLab flybody fruit fly under real MuJoCo physics. No server — everything
+runs in the tab.**
+
+Two separate things move the body, and the difference matters:
+
+- A **trained RL walking policy** (Vaxenburg et al. 2025) genuinely walks the
+  fly from leg actuation and ground reaction alone. That path bypasses the
+  brain and the spine entirely.
+- The **connectome** drives a hand-written tripod gait. It scales that gait but
+  does not generate its rhythm — the leg phase is `sin(sim_time · freq)`.
+
+**The connectome does not walk the body.** The brain→spine link is a cell-type
+name join across two different animals' connectomes, not a reconstructed
+synaptic bridge. Simulated leg-motor pools in this connectome do oscillate
+(5–12 Hz), but in global synchrony rather than in a tripod gait, and no global
+gain setting produces tripod alternation — see `tools/vnc_rhythm.py`.
+
+The brain LIF kernel runs at roughly 0.25 kHz of biological time on an M2 Pro,
+about 4× slower than real time. NEST and Brian2 are faster and better
+validated; the contribution here is reachability, not throughput.
+
+## Running it
+
+Requires **WebGPU** — Chrome, Edge, or Safari 26+. Firefox will not work yet.
+
+About **314 MB** of connectome and body data downloads on first visit and is
+then cached in IndexedDB, so later visits are fast.
+
+For the best experience open the app on its own origin rather than inside the
+Space frame — an embedded page gets a partitioned (or blocked) storage bucket,
+so the 314 MB cache may not survive between visits.
+
+## Credits and licensing
+
+Code is MIT. The data is not, and each piece keeps its own terms:
+
+| | |
+|---|---|
+| Brain connectome | [FlyWire](https://flywire.ai) FAFB, CC-BY |
+| Ventral nerve cord | [Janelia MANC](https://www.janelia.org/project-team/flyem/manc-connectome) (Takemura et al. 2024) |
+| Body model + walking policy | [TuragaLab/flybody](https://github.com/TuragaLab/flybody) (Vaxenburg et al. 2025), Apache-2.0 |
+| Physics | MuJoCo compiled to WebAssembly |
+
+Full attribution in `NOTICE`; every approximation and shortcut is inventoried
+in `LIMITATIONS.md`.
+
+Source: <https://github.com/abgnydn/webgpu-fly>

src/brain.ts

@@ -33,13 +33,16 @@ export interface Brain {
   weight: Float32Array;     // length E
 }
 
-export async function loadBrain(url: string = "/brain.bin"): Promise<Brain> {
+export async function loadBrain(
+  url: string = "/brain.bin",
+  onProgress?: (got: number, total: number) => void,
+): Promise<Brain> {
   // Route through IDB cache so reloads in the same browser session
   // skip the ~125 MB network fetch (~30s on the dev server). Cache
   // key is the full URL including ?v=<sha> from assets.json, so a
   // new build naturally invalidates the cache (different URL = new
   // entry), and idbPut deletes the previous generation of the same asset.
-  const buf = await getOrFetch(url, url);
+  const buf = await getOrFetch(url, url, onProgress);
   return parseBrain(buf);
 }
 

src/cache.ts

@@ -63,12 +63,68 @@ async function idbPut(key: string, value: Entry): Promise<void> {
   });
 }
 
+/** Emit progress at most once per MB — see readWithProgress. */
+const PROGRESS_STEP = 1_000_000;
+
+/** "42.0 / 125.7 MB (33%)", or just "42.0 MB" when the length is unknown. */
+export function progressText(got: number, total: number): string {
+  const mb = (b: number) => (b / 1e6).toFixed(1);
+  return total > 0
+    ? `${mb(got)} / ${mb(total)} MB (${Math.round((100 * got) / total)}%)`
+    : `${mb(got)} MB`;
+}
+
+/**
+ * Drain `r.body` chunk by chunk so the caller can paint a progress line.
+ * Chunks are kept as-is and copied once into a pre-sized buffer — growing
+ * a concatenated array per chunk would be quadratic on a 140 MB bundle.
+ * The callback is throttled to one call per MB: at chunk granularity it
+ * fires thousands of times a second and the DOM writes on the other end
+ * measurably slow the download.
+ */
+async function readWithProgress(
+  r: Response,
+  onProgress: (got: number, total: number) => void,
+): Promise<ArrayBuffer> {
+  const len = Number(r.headers.get("Content-Length"));
+  // Absent / non-numeric / 0 → report bytes only, never "/ 0 MB".
+  const total = Number.isFinite(len) && len > 0 ? len : 0;
+  const reader = r.body!.getReader();
+  const chunks: Uint8Array[] = [];
+  let got = 0;
+  let nextEmit = PROGRESS_STEP;
+  for (;;) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    chunks.push(value);
+    got += value.byteLength;
+    if (got >= nextEmit) {
+      onProgress(got, total);
+      nextEmit = got + PROGRESS_STEP;
+    }
+  }
+  const out = new Uint8Array(got);
+  let off = 0;
+  for (const c of chunks) {
+    out.set(c, off);
+    off += c.byteLength;
+  }
+  onProgress(got, total);
+  return out.buffer;
+}
+
 /**
  * Get bytes for `key` (the cache key) by fetching `url` if not cached.
  * If `key` is already in IDB, return its bytes immediately. Otherwise
  * fetch over network, store, return. Network errors propagate.
+ * `onProgress` (bytes so far, total or 0) streams the download instead of
+ * buffering it whole; without it the response is read in one shot.
  */
-export async function getOrFetch(key: string, url: string): Promise<ArrayBuffer> {
+export async function getOrFetch(
+  key: string,
+  url: string,
+  onProgress?: (got: number, total: number) => void,
+): Promise<ArrayBuffer> {
   try {
     const hit = await idbGet(key);
     if (hit) return hit.bytes;
@@ -78,7 +134,9 @@ export async function getOrFetch(key: string, url: string): Promise<ArrayBuffer>
   }
   const r = await fetch(url);
   if (!r.ok) throw new Error(`${url}: HTTP ${r.status}`);
-  const bytes = await r.arrayBuffer();
+  const bytes = onProgress && r.body
+    ? await readWithProgress(r, onProgress)
+    : await r.arrayBuffer();
   // Fire-and-forget the IDB write — we have the bytes in memory and
   // the caller doesn't need to block on the cache populating. For
   // 125 MB blobs the IDB write is ~30 s and was dominating cold-load

src/main.ts

@@ -10,6 +10,7 @@ import { GaitEvolver } from "./evolution";
 import { loadWalkingPolicy, loadWalkingObsNorm, type WalkingPolicy, type ObsNormStats } from "./walking-policy";
 import { loadWalkingRef } from "./walking-ref";
 import { loadManifest, type VersionFor } from "./manifest";
+import { progressText } from "./cache";
 import { Game, type DnEntry } from "./game";
 
 // ?mode=game (default) shows the playable HUD; ?mode=science keeps the
@@ -137,15 +138,53 @@ function bootSkip(name: "brain" | "vnc" | "body", detail: string) {
 }
 // Surface a fatal boot error on the overlay itself — in game mode the log
 // pane is hidden, so log() alone leaves the overlay spinning forever.
-function bootFail(msg: string) {
+function bootFail(msg: string, link?: { href: string; text: string }) {
   const el = document.querySelector<HTMLElement>("#boot .blink");
   if (!el) return;              // overlay already dismissed/removed
   el.textContent = msg;
   el.style.color = "#ff6b6b";
   el.style.animation = "none";
+  if (link) {
+    const a = document.createElement("a");
+    a.href = link.href;
+    a.textContent = link.text;
+    a.style.cssText = "display: block; margin-top: 10px; color: #9ad7ff";
+    el.appendChild(a);
+  }
 }
 
 async function main() {
+  // Bail before the ~300 MB of binaries, not after: a Safari/Firefox
+  // visitor used to sit through the whole connectome download only to be
+  // told their browser can't run it.
+  if (!("gpu" in navigator)) {
+    log("navigator.gpu missing — needs Chrome, Edge, or Safari 26+", "err");
+    bootFail(
+      "WebGPU unavailable in this browser — the simulator needs Chrome, Edge, or Safari 26+ (Firefox: Windows only, 141+).",
+      // Relative, not root-absolute: the site may be served under a
+      // subpath (Hugging Face Space), where "/index.html" 404s.
+      { href: "index.html", text: "→ read what this project is (no WebGPU needed)" },
+    );
+    return;
+  }
+
+  // Embedded (e.g. the Hugging Face Space page frames the app): storage is
+  // partitioned per-embedder and may be blocked outright, so the ~300 MB
+  // IndexedDB cache can silently fail to persist between visits. Point at
+  // the top-level origin, where it survives.
+  if (window.self !== window.top) {
+    const hint = document.querySelector<HTMLElement>("#boot .blink");
+    if (hint) {
+      const a = document.createElement("a");
+      a.href = window.location.href;
+      a.target = "_blank";
+      a.rel = "noopener";
+      a.textContent = "↗ open in its own tab so the download is cached";
+      a.style.cssText = "display: block; margin-top: 8px; color: #9ad7ff; font-size: 12px";
+      hint.appendChild(a);
+    }
+  }
+
   // Cache-bust manifest. Maps asset filename → "?v=<12-char sha>" so the
   // browser cache key is unique per build. Combined with `Cache-Control:
   // immutable` on R2, subsequent loads hit the browser cache and skip
@@ -169,7 +208,9 @@ async function main() {
   bootStage("brain", "run", "fetching connectome (120 MB)…");
   let brain: Brain;
   try {
-    brain = await loadBrain(brainUrl + versionFor("brain.bin"));
+    brain = await loadBrain(brainUrl + versionFor("brain.bin"), (got, total) => {
+      bootStage("brain", "run", `fetching connectome — ${progressText(got, total)}`);
+    });
   } catch (e) {
     log(`failed: ${(e as Error).message}`, "err");
     log("did you run `npm run data && npm run convert`?", "warn");
@@ -255,7 +296,9 @@ async function main() {
     physicsReject = rej;
   });
   Physics.create((msg) => {
-    log(`flybody: ${msg}`);
+    // Per-megabyte bundle progress belongs on the boot overlay only —
+    // logging it too appends ~140 lines and buries the real trace.
+    if (!msg.startsWith("fetching flybody bundle —")) log(`flybody: ${msg}`);
     bootStage("body", "run", msg);
   })
     .then(async (p) => {
@@ -560,11 +603,6 @@ async function main() {
     driveReadout.textContent = `fwd ${driveFwd.toFixed(2)}  turn ${driveTurn.toFixed(2)}`;
   }
 
-  if (!("gpu" in navigator)) {
-    log("navigator.gpu missing — open in Chrome / Edge", "err");
-    bootFail("WebGPU unavailable — open in Chrome or Edge");
-    return;
-  }
   const sim = await FlySim.create(brain, { ...DEFAULT_PARAMS });
   log(`FlySim ready. dt=${sim.params.dtMs} ms  tau=${sim.params.tauMs} ms`);
 
@@ -588,7 +626,9 @@ async function main() {
   const vncMetaUrl = import.meta.env.VITE_VNC_META_URL || "/vnc.meta.json";
   bootStage("vnc", "run", "fetching MANC connectome (43 MB)…");
   try {
-    const vncBrain = await loadBrain(vncUrl + versionFor("vnc.bin"));
+    const vncBrain = await loadBrain(vncUrl + versionFor("vnc.bin"), (got, total) => {
+      bootStage("vnc", "run", `fetching MANC connectome — ${progressText(got, total)}`);
+    });
     const vncMetaResp = await fetch(vncMetaUrl + versionFor("vnc.meta.json"));
     vncMeta = await vncMetaResp.json();
     vncSim = await FlySim.create(vncBrain, { ...DEFAULT_PARAMS });

src/physics.ts

@@ -12,7 +12,7 @@ import type {
   MainModule, MjModel, MjData,
   MjvScene, MjvOption, MjvPerturb, MjvCamera,
 } from "@mujoco/mujoco";
-import { getOrFetch } from "./cache";
+import { getOrFetch, progressText } from "./cache";
 
 /** Number of actuators that {@link patchActuatorFilters} gives an
  * activation state — 70 `general` + 8 `adhesion` in fruitfly.xml. */
@@ -108,7 +108,9 @@ export class Physics {
 
     onProgress?.("fetching flybody bundle (140 MB, one IDB transaction)");
     const t0 = performance.now();
-    const bundle = await getOrFetch(bundleUrl, bundleUrl);
+    const bundle = await getOrFetch(bundleUrl, bundleUrl, (got, total) => {
+      onProgress?.(`fetching flybody bundle — ${progressText(got, total)}`);
+    });
     onProgress?.(`bundle fetched in ${((performance.now() - t0) / 1000).toFixed(1)} s`);
 
     // Parse bundle: tiny header + manifest, then names + data sections.