Faster WebGPU sampling: fp16 compute, GPU-resident steps, non-blocking thinker preview

- 002/003 and 001 load the v2 build from webgpu/v2/ in each model repo: one graph run per Euler step (CFG mix +
update inside the graph), z stays on the GPU between steps, fp16 compute (fp32 for graph I/O, CFG/Euler,
timestep/resolution/count embeddings and Range), 1x1 conv + depth-to-space instead of ConvTranspose.
- 003 @ 512: 275 -> 161 ms/step, 14 -> 8.1 s/image; 256 px and 001/002: ~6 -> ~3.4 s (RTX 4060, Edge).
- Show thinker: async readback, colouring in a Web Worker, frames dropped while busy; no cost to sampling.
- "Fast" toggle: 25 steps (~4.1 s at 512); default stays 50.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018EtDckYfNgMbkwqVmTsdha

README.md

@@ -37,14 +37,14 @@ this browser.
 | [Preview 001](https://huggingface.co/Logolabs/agate-preview-001) | 256 px | the first preview | 530 MB |
 
 Each model is downloaded once and kept in the browser's Cache Storage. The files live in each release's model repo
-under `webgpu/` (e.g. [Logolabs/agate-preview-003/webgpu](https://huggingface.co/Logolabs/agate-preview-003/tree/main/webgpu)),
+under `webgpu/v2/` (e.g. [Logolabs/agate-preview-003/webgpu/v2](https://huggingface.co/Logolabs/agate-preview-003/tree/main/webgpu/v2)),
 with a `manifest.json` holding sizes and sha256s (the cache keys). `models/` in this Space is the original 001 build
 (without the plan output), kept for `?models=./models/`.
 
 ## Show thinker
 
 Agate's thinker lays the picture out on a 16 × 16 grid (the plan) that steers the renderer. With **Show thinker**
-(on by default; it costs about 2-4% per step) the page shows, at every step, the plan of the conditional branch as
+(on by default; it does not slow sampling: previews are read back asynchronously and coloured in a Web Worker) the page shows, at every step, the plan of the conditional branch as
 colours and the image the model currently expects (x₁ = z + (1 − t)·v, with the SD latent→RGB approximation) —
 the same live preview as the [ComfyUI nodes](https://github.com/logolabs/agate-comfyui): the plan's cells are
 projected on its top three principal components, fitted at the first step with the 2% / 98% quantiles fixed then,
@@ -55,18 +55,24 @@ so a colour keeps its meaning while the plan evolves.
 A browser with WebGPU: a recent Chrome or Edge (113+). Safari and Firefox may lack WebGPU or need it enabled in
 their settings; without it the page offers a WebAssembly (CPU) path that works but takes minutes per image.
 
-## Speed and accuracy (RTX 4060 laptop, Edge, WebGPU, warm)
+## Speed and accuracy (RTX 4060 laptop, Edge, WebGPU, warm, 50 steps)
 
-| model | per step (CFG batch of 2) | 50-step image | model setup per visit |
+| model | per step (CFG batch of 2) | image | before (2026-09-29 morning) |
 |---|---|---|---|
-| 003 at 512 px | ~275 ms | ~14 s | ~15 s |
-| 003 at 256 px | ~120 ms | ~6 s | ~15 s |
-| 002 / 001 | ~105 ms | ~5.3 s | ~14 s |
-
-With identical initial noise, the browser output matches the PyTorch packages (TAESD decoder): 001 within
-5/255 per pixel, 002 within 14/255 (PSNR 52.8 dB), 003 within 29/255 at 512 px (PSNR 54.8 dB) and 47/255 at
-256 px (PSNR 48.9 dB), on a few pixels only (mean difference 0.1-0.2/255). The difference comes from storing the
-weights as fp16 to halve the download; an fp32 build matches to 1/255. Run the self-test with `?parity=1`.
+| 003 at 512 px | ~161 ms | **~8.1 s** (Fast, 25 steps: ~4.1 s) | 275 ms, 14 s |
+| 003 at 256 px | ~70 ms | ~3.5 s | 120 ms, 6.1 s |
+| 002 / 001 | ~65 ms | ~3.3 s | 105 ms, 5.3 s |
+
+How: each step is one graph run that also does the CFG mix and the Euler update, so the latent never leaves the
+GPU between steps (the only readback is the final latent); the network computes in **fp16** (WebGPU `shader-f16`;
+WGSL has no bf16), with fp32 kept for the graph I/O, the CFG/Euler arithmetic, the timestep / resolution / count
+embeddings (sin(1000·t·f) needs fp32) and `Range`; the transposed-conv upsamplers run as 1x1 conv + depth-to-space.
+The thinker preview never blocks sampling: its tensors are downloaded asynchronously, coloured in a Web Worker,
+and frames are dropped while one is pending (measured cost: none).
+
+fp16 compute moves the result slightly: against the fp32 PyTorch reference with the same initial noise, 001 is
+within 16/255 (PSNR 48 dB), 002 36 dB, 003 at 256 px 42 dB and 003 at 512 px 31 dB (mean 0.6-1.8/255). For scale,
+the Python package's own CUDA path (bf16) is 14 dB from that fp32 reference for the same noise. Run `?parity=1`.
 Seeds use a JavaScript PRNG, so a seed gives the same image in every browser, but not the same image as the
 PyTorch pipeline with that seed. The Python packages decode with the full SD-VAE by default; this page uses TAESD.
 

css/style.css

@@ -130,3 +130,5 @@ input[type=text] { width: 100%; font: 15px/1.45 var(--sans); color: var(--ink);
 .thinker #pred-canvas { image-rendering: auto; }
 .thinker figcaption { margin-top: 6px; }
 .thinker-note { grid-column: 1 / -1; margin: 0; font-size: 13px; color: var(--mute); }
+.label .toggle.fast { margin-left: 10px; text-transform: uppercase; letter-spacing: .1em; color: var(--ink); }
+.label .toggle.fast input { accent-color: var(--red); margin: 0 4px 0 0; vertical-align: -1px; }

index.html

@@ -68,7 +68,8 @@
           </div>
         </div>
         <div class="field">
-          <label class="label" for="steps">Steps <b id="steps-v">50</b></label>
+          <label class="label" for="steps">Steps <b id="steps-v">50</b>
+            <span class="toggle fast" title="25 steps: about twice as fast; detail holds, the layout can change"><input type="checkbox" id="fast"> Fast</span></label>
           <input id="steps" type="range" min="4" max="50" step="1" value="50">
         </div>
         <div class="field">

js/agate.js

@@ -15,7 +15,7 @@ export { ort };
 // (001 without the plan output) is used only through ?models=. dir: the layout under a local
 // models base (?models=<url> or window.AGATE_MODEL_BASE given explicitly, e.g. for local testing). cache: Cache
 // Storage name (001 keeps the name it always had, so returning visitors do not download it again).
-const HUB = (v) => `https://huggingface.co/Logolabs/agate-preview-${v}/resolve/main/webgpu/`;
+const HUB = (v) => `https://huggingface.co/Logolabs/agate-preview-${v}/resolve/main/webgpu/v2/`;   // v2: fast fp16 step graphs
 export const VERSIONS = {
   "003": { remote: HUB("003"), dir: "003/", cache: "agate-preview-003-web", res: [512, 256], note: "Multi-resolution model, 512 px (or 256). Quoted text is spelled out, counts are coded, “no X” becomes a negative prompt." },
   "002": { remote: HUB("002"), dir: "002/", cache: "agate-preview-002-web", res: [256], note: "Same network as 001, trained 38,710 steps longer before the same anneal. 256 px, plain prompts." },
@@ -99,6 +99,9 @@ export class Agate {
 
   get mr() { return this.manifest?.arch === "fcdm_t2mr"; }
 
+  // the fast build (2026-09-29): one Euler step per run, fp16 compute, everything stays on the GPU between steps
+  get fast() { return !!this.manifest?.step; }
+
   fileKey(f) { return new Request(`${this.base}${f}?sha256=${this.manifest.files[f].sha256}`); }
 
   async getFile(f, onBytes = () => {}) {
@@ -115,7 +118,7 @@ export class Agate {
     this.resolution = this.mr ? Number(resolution || manifest.resolution) : 256;
     // 003: download everything (the two graphs are ~1 MB each and share one weights file), so a resolution
     // switch needs no network
-    const need = this.mr ? Object.keys(manifest.files) : FILES_T2;
+    const need = this.mr || manifest.step ? Object.keys(manifest.files) : FILES_T2;
     const total = need.reduce((s, f) => s + manifest.files[real(f)].bytes, 0);
     let loaded = 0, fromCache = 0;
     const bufs = {};
@@ -151,7 +154,10 @@ export class Agate {
     }
     onProgress({ phase: "init", file: "generator", loaded: total, total });
     const ts = performance.now();
-    if (this.mr) {
+    if (this.fast) {
+      this.weights = bufs[manifest.step.external_data];      // kept for a resolution switch (no re-read)
+      await this.createGenerator(this.resolution, bufs);
+    } else if (this.mr) {
       this.weights = bufs["generator.weights"];              // kept for a resolution switch (no re-read)
       await this.createGenerator(this.resolution, bufs);
     } else {
@@ -166,12 +172,20 @@ export class Agate {
 
   // 003: the generator graph for one resolution; the weights file is shared by both graphs.
   async createGenerator(res, bufs = {}) {
-    const r = this.manifest.resolutions[String(res)];
-    if (!r) throw new Error(`resolution ${res} not available`);
-    const graph = bufs[r.graph] || await this.getFile(r.graph);
-    const weights = this.weights || await this.getFile(r.external_data);
+    let graphName, ext, extra = {};
+    if (this.fast) {
+      graphName = this.manifest.step.graphs[String(res)];
+      ext = this.manifest.step.external_data;
+      extra = { preferredOutputLocation: "gpu-buffer" };     // z_next / x1 / plan stay on the GPU
+    } else {
+      const r = this.manifest.resolutions[String(res)];
+      graphName = r?.graph; ext = r?.external_data;
+    }
+    if (!graphName) throw new Error(`resolution ${res} not available`);
+    const graph = bufs[graphName] || await this.getFile(graphName);
+    const weights = this.weights || await this.getFile(ext);
     if (this.sessions.gen) { try { await this.sessions.gen.release(); } catch { /* */ } this.sessions.gen = null; }
-    this.sessions.gen = await ort.InferenceSession.create(graph, this.sessionOptions({ externalData: [{ path: r.external_data, data: weights }] }));
+    this.sessions.gen = await ort.InferenceSession.create(graph, this.sessionOptions({ externalData: [{ path: ext, data: weights }], ...extra }));
     this.resolution = Number(res);
   }
 
@@ -240,9 +254,57 @@ export class Agate {
   // -> { rgba (watermarked unless watermark=false), width, height, latent, timings, prepared }
   get hasPlan() { return !!this.sessions?.gen?.outputNames?.includes("plan"); }
 
+  // GPU-resident sampler for the step graphs: z_next of one run is the z input of the next, ctx / mask / counts
+  // are uploaded once, and nothing is read back until the end. The preview (plan + x1) is requested only when
+  // onPreview is set, previewReady() says the viewer is free and no earlier readback is still in flight; its
+  // download is started but never awaited by the loop (frames are dropped instead).
+  async sampleFast({ z, grid, cfg, hw, C, L, D, ctx, mask, cnt, onStep, onPreview, previewReady, shouldStop, stepMs }) {
+    const dev = ort.env.webgpu?.device;
+    const bufs = [];
+    const gpuT = (data, dims) => {
+      if (!dev) return new ort.Tensor("float32", data, dims);
+      const buf = dev.createBuffer({ size: Math.ceil(data.byteLength / 16) * 16, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
+      dev.queue.writeBuffer(buf, 0, data);
+      bufs.push(buf);
+      return ort.Tensor.fromGpuBuffer(buf, { dataType: "float32", dims });
+    };
+    const fixed = { ctx: gpuT(ctx, [2, L, D]), mask: gpuT(mask, [2, L]) };
+    if (this.sessions.gen.inputNames.includes("counts")) fixed.counts = gpuT(cnt || new Float32Array(2 * L), [2, L]);
+    const steps = grid.length - 1;
+    let zT = new ort.Tensor("float32", z, [1, C, hw, hw]), inflight = false;
+    const one = (v) => new ort.Tensor("float32", new Float32Array([v]), [1]);
+    const cfgT = one(cfg);
+    try {
+      for (let i = 0; i < steps; i++) {
+        if (shouldStop()) throw new Error("stopped");
+        const ts = performance.now();
+        const want = !!onPreview && this.hasPlan && !inflight && previewReady();
+        const out = await this.sessions.gen.run({ z: zT, t: one(grid[i]), dt: one(grid[i + 1] - grid[i]), cfg: cfgT, ...fixed },
+                                                want ? ["z_next", "x1", "plan"] : ["z_next"]);
+        if (i > 0) zT.dispose?.();
+        zT = out.z_next;
+        if (want) {
+          inflight = true;
+          const step = i + 1;
+          Promise.all([out.plan.getData(true), out.x1.getData(true)])
+            .then(([plan, x1]) => onPreview({ step, steps, plan, x1, hw }))
+            .catch((e) => console.warn("preview", e))
+            .finally(() => { inflight = false; });
+        }
+        stepMs.push(performance.now() - ts);
+        onStep(i + 1, steps, null);
+        if ((i & 3) === 3) await new Promise((r) => setTimeout(r, 0));   // let the page paint now and then
+      }
+      const zf = await zT.getData(true);                    // the only synchronising readback
+      return Float32Array.from(zf);
+    } finally {
+      for (const b of bufs) { try { b.destroy(); } catch { /* */ } }
+    }
+  }
+
   // onPreview({ step, steps, plan, x1, hw }) after every step, when given and the graph outputs the plan: plan =
   // the thinker output for the conditional branch (Float32Array 640 x 16 x 16), x1 = z + (1 - t) v (guided).
-  async generate({ prompt, negative = "", seed = 0, steps = 50, cfg = 3.0, noise = null, watermark = true, onStep = () => {}, onPreview = null, shouldStop = () => false }) {
+  async generate({ prompt, negative = "", seed = 0, steps = 50, cfg = 3.0, noise = null, watermark = true, onStep = () => {}, onPreview = null, previewReady = () => true, shouldStop = () => false }) {
     const C = this.manifest.latent_ch, D = this.manifest.ctx_dim;
     const rc = this.mr ? this.manifest.resolutions[String(this.resolution)] : { latent_hw: this.manifest.latent_hw, shift: 1 };
     const hw = rc.latent_hw, shift = Number(rc.shift || 1);
@@ -268,10 +330,13 @@ export class Agate {
     }
 
     let z = noise ? Float32Array.from(noise) : Agate.gaussianNoise(seed, per);
-    const zz = new Float32Array(2 * per), tt = new Float32Array(2);
     const grid = Array.from({ length: steps + 1 }, (_, i) => Agate.shiftT(i / steps, shift));
     t0 = performance.now();
     const stepMs = [];
+    if (this.fast) {
+      z = await this.sampleFast({ z, grid, cfg, hw, C, L, D, ctx, mask, cnt: feeds.counts?.data, onStep, onPreview, previewReady, shouldStop, stepMs });
+    } else {
+    const zz = new Float32Array(2 * per), tt = new Float32Array(2);
     for (let i = 0; i < steps; i++) {
       if (shouldStop()) throw new Error("stopped");
       const ts = performance.now();
@@ -296,6 +361,7 @@ export class Agate {
       onStep(i + 1, steps, z);
       await new Promise((r) => setTimeout(r, 0));          // let the progress bar paint
     }
+    }
     T.samplerMs = performance.now() - t0;
     T.firstStepMs = stepMs[0];
     T.stepMs = stepMs.length > 1 ? (T.samplerMs - stepMs[0]) / (stepMs.length - 1) : stepMs[0];

js/app.js

@@ -116,7 +116,8 @@ async function pickRes(r) {
   }
 }
 
-ui.steps.oninput = () => { $("steps-v").textContent = ui.steps.value; };
+ui.steps.oninput = () => { $("steps-v").textContent = ui.steps.value; $("fast").checked = Number(ui.steps.value) === 25; };
+$("fast").onchange = () => { ui.steps.value = $("fast").checked ? 25 : 50; $("steps-v").textContent = ui.steps.value; };
 ui.cfg.oninput = () => { $("cfg-v").textContent = Number(ui.cfg.value).toFixed(1); };
 ui.dice.onclick = () => { ui.seed.value = Math.floor(Math.random() * 2 ** 31); };
 ui.crisp.checked = store.get("agate-crisp") === "1";
@@ -169,18 +170,19 @@ async function pngWithMetadata() {
   return new Blob([bytes], { type: "image/png" });
 }
 
+// The previews: frames go to thinkerView, which colours them in a Web Worker and drops frames while busy. The
+// sampler never waits for any of it (agate.js sampleFast); with Show thinker off nothing is requested at all.
 function previewFn() {
   thinkerView.reset();
-  let planMs = 0, n = 0;
-  window.__agate.thinker = { steps: 0, planMs: 0 };
-  return ({ step, steps, plan, x1, hw }) => {
-    const t0 = performance.now();
+  let n = 0, t0 = performance.now();
+  window.__agate.thinker = { submitted: 0 };
+  thinkerView.onDrawn = ({ step, steps, workerMs, dropped, drawn }) => {
     ui.thinker.hidden = false;
-    thinkerView.drawPlan(plan);
-    thinkerView.drawPrediction(x1, hw);
     ui.thinkerStep.textContent = `${step} / ${steps}`;
-    planMs += performance.now() - t0; n++;
-    window.__agate.thinker = { steps: n, drawMsMean: planMs / n };
+    window.__agate.thinker = { submitted: n, drawn, dropped, lastWorkerMs: workerMs, lastStep: step };
+  };
+  return ({ step, steps, plan, x1, hw }) => {
+    if (thinkerView.submit(plan, x1, hw, { step, steps })) n++;
   };
 }
 
@@ -197,6 +199,7 @@ async function generate(opts = {}) {
       watermark: opts.watermark ?? true,
       onStep: (i, n) => setProgress(i / n, `Step ${i} / ${n}`, ""),
       onPreview: (opts.thinker ?? ui.showThinker.checked) ? previewFn() : null,
+      previewReady: () => !thinkerView.pending,
       shouldStop: () => stop,
     });
     const ctx = ui.canvas.getContext("2d");

js/thinker-worker.js

@@ -0,0 +1,11 @@
+// Web Worker: colours the thinker previews off the main thread (see thinker.js).
+import { planPixels, predPixels } from "./thinker.js";
+
+let gen = -1, state = {};
+self.onmessage = (e) => {
+  const { gen: g, plan, x1, hw, meta } = e.data;
+  if (g !== gen) { gen = g; state = {}; }           // a new generation: fit a new basis at its first frame
+  const t0 = performance.now();
+  const p = planPixels(plan, state), q = predPixels(x1, hw);
+  self.postMessage({ gen: g, plan: p, pred: q, hw, meta, ms: performance.now() - t0 }, [p.buffer, q.buffer]);
+};

js/thinker.js

@@ -75,42 +75,75 @@ function quantile(sorted, q) {                 // torch.quantile, linear interpo
   return sorted[lo] + (sorted[hi] - sorted[lo]) * (pos - lo);
 }
 
+// state: { basis, lo, hi } (filled at the first call of a generation). -> RGBA Uint8ClampedArray (g * g * 4)
+export function planPixels(plan, state, C = 640, g = 16) {
+  const n = g * g, X = new Float64Array(n * C);
+  for (let c = 0; c < C; c++) for (let p = 0; p < n; p++) X[p * C + c] = plan[c * n + p];
+  for (let c = 0; c < C; c++) { let m = 0; for (let p = 0; p < n; p++) m += X[p * C + c]; m /= n; for (let p = 0; p < n; p++) X[p * C + c] -= m; }
+  if (!state.basis) state.basis = topComponents(X, n, C);
+  const Y = state.basis.map((v) => { const y = new Float64Array(n); for (let p = 0; p < n; p++) { let s = 0; const r = p * C; for (let c = 0; c < C; c++) s += X[r + c] * v[c]; y[p] = s; } return y; });
+  if (!state.lo) {
+    state.lo = Y.map((y) => quantile(Float64Array.from(y).sort(), 0.02));
+    state.hi = Y.map((y) => quantile(Float64Array.from(y).sort(), 0.98));
+  }
+  const px = new Uint8ClampedArray(n * 4);
+  for (let p = 0; p < n; p++) {
+    for (let k = 0; k < 3; k++) px[p * 4 + k] = Math.round(Math.min(1, Math.max(0, (Y[k][p] - state.lo[k]) / (state.hi[k] - state.lo[k] + 1e-8))) * 255);
+    px[p * 4 + 3] = 255;
+  }
+  return px;
+}
+
+// x1: Float32Array (4 * hw * hw) in the model's (scaled SD-VAE) latent space -> RGBA (hw * hw * 4)
+export function predPixels(x1, hw) {
+  const n = hw * hw, px = new Uint8ClampedArray(n * 4);
+  for (let p = 0; p < n; p++) {
+    for (let k = 0; k < 3; k++) {
+      let s = 0; for (let c = 0; c < 4; c++) s += x1[c * n + p] * RGB_FACTORS[c][k];
+      px[p * 4 + k] = Math.floor(Math.min(1, Math.max(0, (s + 1) / 2)) * 255);
+    }
+    px[p * 4 + 3] = 255;
+  }
+  return px;
+}
+
+// Draws previews; the colouring runs in a Web Worker (thinker-worker.js) when available, so the page and the
+// sampler never wait for it. A frame arriving while the previous one is still being coloured is dropped.
 export class ThinkerView {
-  constructor(planCanvas, predCanvas) { this.planCanvas = planCanvas; this.predCanvas = predCanvas; this.reset(); }
+  constructor(planCanvas, predCanvas, onDrawn = () => {}) {
+    this.planCanvas = planCanvas; this.predCanvas = predCanvas; this.onDrawn = onDrawn;
+    this.gen = 0; this.busy = false; this.dropped = 0; this.drawn = 0;
+    try {
+      this.worker = new Worker(new URL("./thinker-worker.js", import.meta.url), { type: "module" });
+      this.worker.onmessage = (e) => this._paint(e.data);
+    } catch { this.worker = null; }
+    this.reset();
+  }
 
-  reset() { this.basis = null; this.lo = null; this.hi = null; }
+  reset() { this.gen++; this.state = {}; this.busy = false; this.dropped = 0; this.drawn = 0; }
 
-  // plan: Float32Array (C * g * g), channel-major (1, C, g, g)
-  drawPlan(plan, C = 640, g = 16) {
-    const n = g * g, X = new Float64Array(n * C);
-    for (let c = 0; c < C; c++) for (let p = 0; p < n; p++) X[p * C + c] = plan[c * n + p];
-    for (let c = 0; c < C; c++) { let m = 0; for (let p = 0; p < n; p++) m += X[p * C + c]; m /= n; for (let p = 0; p < n; p++) X[p * C + c] -= m; }
-    if (!this.basis) this.basis = topComponents(X, n, C);
-    const Y = this.basis.map((v) => { const y = new Float64Array(n); for (let p = 0; p < n; p++) { let s = 0; const r = p * C; for (let c = 0; c < C; c++) s += X[r + c] * v[c]; y[p] = s; } return y; });
-    if (!this.lo) {
-      this.lo = Y.map((y) => quantile(Float64Array.from(y).sort(), 0.02));
-      this.hi = Y.map((y) => quantile(Float64Array.from(y).sort(), 0.98));
-    }
-    const img = new ImageData(g, g);
-    for (let p = 0; p < n; p++) {
-      for (let k = 0; k < 3; k++) img.data[p * 4 + k] = Math.round(Math.min(1, Math.max(0, (Y[k][p] - this.lo[k]) / (this.hi[k] - this.lo[k] + 1e-8))) * 255);
-      img.data[p * 4 + 3] = 255;
-    }
-    const cv = this.planCanvas; cv.width = g; cv.height = g;
-    cv.getContext("2d").putImageData(img, 0, 0);
+  get pending() { return this.busy; }
+
+  // plan: Float32Array (640 * 16 * 16), x1: Float32Array (4 * hw * hw); both are transferred to the worker
+  submit(plan, x1, hw, meta = {}) {
+    if (this.busy) { this.dropped++; return false; }
+    this.busy = true;
+    const job = { gen: this.gen, plan, x1, hw, meta };
+    if (this.worker) this.worker.postMessage(job, [plan.buffer, x1.buffer]);
+    else setTimeout(() => this._paint({ gen: job.gen, hw, meta, plan: planPixels(plan, this.state), pred: predPixels(x1, hw) }), 0);
+    return true;
   }
 
-  // x1: Float32Array (4 * hw * hw) in the model's (scaled SD-VAE) latent space
-  drawPrediction(x1, hw) {
-    const n = hw * hw, img = new ImageData(hw, hw);
-    for (let p = 0; p < n; p++) {
-      for (let k = 0; k < 3; k++) {
-        let s = 0; for (let c = 0; c < 4; c++) s += x1[c * n + p] * RGB_FACTORS[c][k];
-        img.data[p * 4 + k] = Math.floor(Math.min(1, Math.max(0, (s + 1) / 2)) * 255);
-      }
-      img.data[p * 4 + 3] = 255;
-    }
-    const cv = this.predCanvas; cv.width = hw; cv.height = hw;
-    cv.getContext("2d").putImageData(img, 0, 0);
+  _paint({ gen, plan, pred, hw, meta, ms }) {
+    this.busy = false;
+    if (gen !== this.gen) return;
+    const put = (cv, px, s) => { cv.width = s; cv.height = s; cv.getContext("2d").putImageData(new ImageData(px, s, s), 0, 0); };
+    put(this.planCanvas, plan, 16);
+    put(this.predCanvas, pred, hw);
+    this.drawn++;
+    this.onDrawn({ ...meta, workerMs: ms, dropped: this.dropped, drawn: this.drawn });
   }
+
+  // synchronous path (tests): colour on this thread
+  drawPlan(plan) { const px = planPixels(plan, this.state); const cv = this.planCanvas; cv.width = 16; cv.height = 16; cv.getContext("2d").putImageData(new ImageData(px, 16, 16), 0, 0); }
 }