Agate 4-step WebGPU: in-browser runtime with live mode

README.md

@@ -1,10 +1,40 @@
 ---
-title: Agate 4step Webgpu
-emoji: 🐨
-colorFrom: blue
-colorTo: yellow
+title: Agate 4-step WebGPU
+emoji: ⚡
+colorFrom: red
+colorTo: gray
 sdk: static
+app_file: index.html
 pinned: false
+license: mit
+short_description: 4-step Agate text-to-image in your browser, live as you type
+custom_headers:
+  cross-origin-embedder-policy: require-corp
+  cross-origin-opener-policy: same-origin
+  cross-origin-resource-policy: cross-origin
+models:
+- ysharma/agate-002-4step
+- Logolabs/agate-preview-002
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Agate 4-step, in your browser
+
+[ysharma/agate-002-4step](https://huggingface.co/ysharma/agate-002-4step) is an **unofficial** 4-step distillation of
+[Agate Preview 002](https://huggingface.co/Logolabs/agate-preview-002) by LogoLabs. This page runs it entirely on your
+GPU with WebGPU (onnxruntime-web). Prompts never leave your machine.
+
+- **4 network passes per image** instead of Agate's 100: guidance is baked into the weights, and the sampler is 4 uniform Euler steps.
+- **Live mode** redraws the image on every keystroke. Keystrokes typed while an image is being made are coalesced, so only the newest prompt gets drawn next.
+- The model files (about 615 MB) download once from the model repo's `webgpu/` folder and are then kept in the browser's Cache Storage.
+
+The runtime (tokenizer, text encoder, sampler, TAESD decoder, watermark) is adapted from LogoLabs'
+[agate-webgpu](https://huggingface.co/spaces/Logolabs/agate-webgpu) (MIT). The one change to the sampler: with `cfg = 1`
+the generator runs a batch of 1, not a conditional/unconditional pair.
+
+**Limitations.** 256 × 256 only; no negative prompts; Agate's weaknesses (exact text, counts above three, negation) remain.
+No safety filter runs in this page, and the base model can draw people partly unclothed without being asked.
+Every image carries Agate's invisible watermark (`AGATE002`), and saved PNGs include AI-generated provenance fields.
+
+Query parameters: `?ep=wasm` forces the CPU backend, `?autoload=1` loads the model on open.
+
+Not affiliated with LogoLabs. Licence: MIT.

css/style.css

@@ -0,0 +1,79 @@
+:root {
+  --bg: #f7f6f3;
+  --surface: #ffffff;
+  --text: #1d1c1a;
+  --muted: #6b6862;
+  --line: #e2dfd8;
+  --accent: #b4462f;
+  --accent-text: #ffffff;
+  --frame: #ecE9e2;
+}
+@media (prefers-color-scheme: dark) {
+  :root:not([data-theme="light"]) {
+    --bg: #151412;
+    --surface: #1f1d1a;
+    --text: #efece6;
+    --muted: #a39f97;
+    --line: #34312c;
+    --accent: #e0765d;
+    --accent-text: #1a0f0c;
+    --frame: #26231f;
+  }
+}
+:root[data-theme="dark"] {
+  --bg: #151412; --surface: #1f1d1a; --text: #efece6; --muted: #a39f97;
+  --line: #34312c; --accent: #e0765d; --accent-text: #1a0f0c; --frame: #26231f;
+}
+* { box-sizing: border-box; }
+html, body { margin: 0; background: var(--bg); color: var(--text); }
+body { font: 15px/1.55 Inter, system-ui, sans-serif; }
+a { color: var(--accent); }
+.page { max-width: 880px; margin: 0 auto; padding: 32px 16px 64px; }
+h1 { font-size: 28px; line-height: 1.2; margin: 0 0 8px; font-weight: 700; letter-spacing: -0.01em; }
+h2 { font-size: 17px; margin: 0 0 8px; }
+.tag { font: 500 12px/1 "JetBrains Mono", monospace; vertical-align: middle; padding: 4px 8px; border-radius: 999px;
+  border: 1px solid var(--line); color: var(--muted); }
+.lede { color: var(--muted); margin: 0 0 24px; max-width: 64ch; }
+.panel, .result, .about, .notice { background: var(--surface); border: 1px solid var(--line); border-radius: 12px; padding: 16px; margin-bottom: 16px; }
+.notice { border-color: var(--accent); }
+.label { display: block; font-weight: 600; font-size: 13px; margin: 12px 0 6px; }
+textarea { width: 100%; resize: vertical; font: inherit; color: var(--text); background: var(--bg);
+  border: 1px solid var(--line); border-radius: 8px; padding: 10px 12px; }
+textarea:focus, input:focus { outline: 2px solid var(--accent); outline-offset: 1px; }
+button { font: 600 14px/1 Inter, system-ui, sans-serif; border-radius: 8px; padding: 10px 14px; cursor: pointer; }
+button:disabled { opacity: 0.5; cursor: default; }
+.primary { background: var(--accent); color: var(--accent-text); border: 1px solid var(--accent); }
+.ghost { background: transparent; color: var(--text); border: 1px solid var(--line); }
+.loadrow { margin-bottom: 10px; }
+.loadrow .primary { width: 100%; }
+.bar { height: 4px; background: var(--line); border-radius: 2px; overflow: hidden; }
+#bar-fill { height: 100%; width: 0; background: var(--accent); transition: width 0.15s; }
+.status { display: flex; justify-content: space-between; gap: 12px; font: 12px/1.4 "JetBrains Mono", monospace; color: var(--muted); margin: 6px 0 0; }
+.chips { display: flex; flex-wrap: wrap; gap: 6px; margin-top: 8px; }
+.chip { font: 12px/1.3 Inter, system-ui, sans-serif; padding: 6px 10px; border-radius: 999px; border: 1px solid var(--line);
+  background: transparent; color: var(--muted); }
+.chip:hover { color: var(--text); border-color: var(--muted); }
+.controls { display: flex; flex-wrap: wrap; align-items: center; gap: 10px; margin-top: 14px; }
+.controls .primary { margin-left: auto; }
+.switch { display: inline-flex; align-items: center; gap: 8px; font-weight: 500; }
+.switch input { width: 16px; height: 16px; accent-color: var(--accent); }
+.seed { display: inline-flex; align-items: center; gap: 6px; color: var(--muted); font-size: 13px; }
+.seed input { width: 120px; font: 13px "JetBrains Mono", monospace; color: var(--text); background: var(--bg);
+  border: 1px solid var(--line); border-radius: 8px; padding: 8px; }
+.result { display: grid; grid-template-columns: minmax(0, 512px) minmax(0, 1fr); gap: 16px; align-items: start; }
+.frame { position: relative; aspect-ratio: 1; background: var(--frame); border-radius: 8px; overflow: hidden; }
+.frame canvas { width: 100%; height: 100%; display: block; image-rendering: auto; }
+.placeholder { position: absolute; inset: 0; display: grid; place-items: center; color: var(--muted); font-size: 14px; }
+.frame.busy canvas { opacity: 0.85; }
+.timings { display: grid; grid-template-columns: auto auto; gap: 4px 12px; margin: 0 0 12px; font: 12px/1.5 "JetBrains Mono", monospace; }
+.timings dt { color: var(--muted); }
+.timings dd { margin: 0; text-align: right; font-variant-numeric: tabular-nums; }
+.timings .big { font-size: 22px; color: var(--text); font-weight: 500; }
+.ai { font-size: 12px; color: var(--muted); }
+.about p { color: var(--muted); max-width: 72ch; }
+.fine { font-size: 12px; }
+@media (max-width: 640px) {
+  .result { grid-template-columns: 1fr; }
+  .controls .primary { margin-left: 0; width: 100%; }
+  h1 { font-size: 24px; }
+}

index.html

@@ -1,19 +1,73 @@
 <!doctype html>
-<html>
-	<head>
-		<meta charset="utf-8" />
-		<meta name="viewport" content="width=device-width" />
-		<title>My static Space</title>
-		<link rel="stylesheet" href="style.css" />
-	</head>
-	<body>
-		<div class="card">
-			<h1>Welcome to your static Space!</h1>
-			<p>You can modify this app directly by editing <i>index.html</i> in the Files and versions tab.</p>
-			<p>
-				Also don't forget to check the
-				<a href="https://huggingface.co/docs/hub/spaces" target="_blank">Spaces documentation</a>.
-			</p>
-		</div>
-	</body>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>Agate 4-step WebGPU</title>
+<meta name="description" content="An unofficial 4-step distillation of Agate Preview 002 (LogoLabs) running entirely in your browser on WebGPU. Redraws as you type.">
+<link rel="preconnect" href="https://fonts.googleapis.com">
+<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+<link crossorigin="anonymous" href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap" rel="stylesheet">
+<link rel="stylesheet" href="css/style.css">
+</head>
+<body>
+<main class="page">
+  <header>
+    <h1>Agate 4-step <span class="tag">WebGPU</span></h1>
+    <p class="lede">A 0.19B text-to-image model distilled to <b>4 steps with no guidance pass</b>, running on your own GPU.
+      Nothing leaves your machine. Turn on <b>Live</b> and the image redraws as you type.</p>
+  </header>
+
+  <div id="nogpu" class="notice" hidden>
+    <p id="nogpu-why"></p>
+    <p>This page needs WebGPU (recent Chrome or Edge). You can still run it on the CPU, which takes a few seconds per image.</p>
+    <button id="use-wasm" class="ghost">Use CPU instead</button>
+  </div>
+
+  <section class="panel">
+    <div class="loadrow" id="loadrow">
+      <button id="load" class="primary">Load model (about 615 MB, cached after the first visit)</button>
+    </div>
+    <div class="bar" aria-hidden="true"><div id="bar-fill"></div></div>
+    <p class="status"><span id="status">Not loaded</span><span id="status-r"></span></p>
+
+    <label for="prompt" class="label">Prompt</label>
+    <textarea id="prompt" rows="3" spellcheck="false" placeholder="a red cube on top of a blue sphere">a lighthouse on a rocky coast under a stormy sky</textarea>
+    <div class="chips" id="examples"></div>
+
+    <div class="controls">
+      <label class="switch"><input type="checkbox" id="live" checked> <span>Live (redraw on every keystroke)</span></label>
+      <label class="seed">Seed <input id="seed" type="number" value="0" min="0" max="4294967295"></label>
+      <button id="dice" class="ghost" title="New seed">New seed</button>
+      <button id="go" class="primary" disabled>Generate</button>
+    </div>
+  </section>
+
+  <section class="result">
+    <div class="frame" id="frame">
+      <canvas id="canvas" width="256" height="256"></canvas>
+      <div class="placeholder" id="placeholder">Load the model to start</div>
+    </div>
+    <div class="side">
+      <dl class="timings" id="timings"></dl>
+      <button id="save" class="ghost" disabled>Save PNG</button>
+      <p class="ai" id="ai-note" hidden>AI-generated image. It carries an invisible watermark and provenance fields in the saved PNG.</p>
+    </div>
+  </section>
+
+  <section class="about">
+    <h2>What this is</h2>
+    <p><a href="https://huggingface.co/ysharma/agate-002-4step" target="_blank" rel="noopener">ysharma/agate-002-4step</a> is an
+      <b>unofficial</b> distillation of <a href="https://huggingface.co/Logolabs/agate-preview-002" target="_blank" rel="noopener">Agate Preview 002</a> by LogoLabs.
+      Classifier-free guidance was baked into the weights, then the 50-step sampler was halved four times (32 → 16 → 8 → 4).
+      Each image is 4 network passes instead of Agate's 100. On the official GenEval scorer it gets 0.527, against 0.562 for the teacher at 50 steps and 0.466 for the teacher run at 4 steps.</p>
+    <p>256 × 256 only. No negative prompts (guidance is baked in). It inherits Agate's weaknesses: exact text, counts above three, negation.
+      There is no safety filter in this page, and the base model can draw people partly unclothed without being asked.</p>
+    <p class="fine">Runtime adapted from LogoLabs' <a href="https://huggingface.co/spaces/Logolabs/agate-webgpu" target="_blank" rel="noopener">agate-webgpu</a> (MIT):
+      tokenizer, text encoder, sampler and TAESD decoder run through onnxruntime-web. Seeds use a JavaScript PRNG and do not match PyTorch seeds.
+      This page is not affiliated with LogoLabs. Licence: MIT.</p>
+  </section>
+</main>
+<script type="module" src="js/app.js"></script>
+</body>
 </html>

js/agate.js

@@ -0,0 +1,390 @@
+// Agate in the browser: tokenizer (tokenizers.js) + ONNX graphs run by onnxruntime-web (WebGPU, WASM
+// fallback) + an Euler flow sampler with CFG. Mirrors agate/pipeline.py (AgatePipeline) of
+// Logolabs/agate-preview-001 / -002 (fcdm_t2, 256 px) and Logolabs/agate-preview-003 (fcdm_t2mr: 512 or 256 px,
+// SD3 timestep shift, prompt pipeline + count code, see prompt.js). Every image is marked (marking.js).
+
+import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.webgpu.bundle.min.mjs";
+import { Tokenizer } from "https://cdn.jsdelivr.net/npm/@huggingface/tokenizers@0.2.0/dist/tokenizers.mjs";
+import { prepare, countVector } from "./prompt.js";
+import { embedWatermark, marks } from "./marking.js";
+export const MODEL_REPO = "ysharma/agate-002-4step";
+
+export { ort };
+
+// The model this page runs: the unofficial 4-step distillation of Agate Preview 002
+// (ysharma/agate-002-4step, MIT). Its webgpu/ folder has the same layout and graph I/O as 002's original webgpu/
+// build (generator.onnx: z, t, ctx, mask -> v, plan; batch and text length dynamic).
+const REPO = MODEL_REPO;
+export const VERSIONS = {
+  "4step": { remote: `https://huggingface.co/${REPO}/resolve/main/webgpu/`, dir: "", cache: "agate-002-4step-web", res: [256], release: "002",
+             note: "Guidance- and progressively-distilled from Agate Preview 002: 4 Euler steps, no CFG. 256 px." },
+};
+const countUrl = () => `https://huggingface.co/${REPO}/resolve/main/config.json`;
+const FILES_T2 = ["tokenizer.json", "tokenizer_config.json", "text_encoder.onnx", "generator.onnx", "taesd_decoder.onnx"];
+
+export async function webgpuStatus() {
+  if (!("gpu" in navigator)) return { ok: false, why: "This browser does not expose WebGPU (navigator.gpu is missing)." };
+  try {
+    const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
+    if (!adapter) return { ok: false, why: "WebGPU is present but no GPU adapter is available." };
+    let name = "";
+    try { const info = adapter.info || (await adapter.requestAdapterInfo?.()); name = [info?.vendor, info?.architecture, info?.description].filter(Boolean).join(" "); } catch { /* optional */ }
+    return { ok: true, adapter: name, maxBuffer: adapter.limits?.maxBufferSize };
+  } catch (e) {
+    return { ok: false, why: `WebGPU adapter request failed: ${e.message || e}` };
+  }
+}
+
+// ---- downloads with progress + Cache Storage ---------------------------------------------------
+async function openCache(name) {
+  try { return await caches.open(name); } catch { return null; }   // file://, private mode, ...
+}
+
+async function fetchCached(cacheName, url, key, expectBytes, onBytes) {
+  const cache = await openCache(cacheName);
+  if (cache) {
+    const hit = await cache.match(key);
+    if (hit) { const buf = new Uint8Array(await hit.arrayBuffer()); onBytes(buf.byteLength, true); return { buf, cached: true }; }
+  }
+  const res = await fetch(url);
+  if (!res.ok) throw new Error(`${url}: HTTP ${res.status}`);
+  // manifest size first: a compressed response would report the encoded length
+  const total = expectBytes || Number(res.headers.get("content-length")) || 0;
+  const reader = res.body.getReader();
+  let buf = new Uint8Array(total || 1 << 20), n = 0;
+  for (;;) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    if (n + value.byteLength > buf.byteLength) { const b2 = new Uint8Array(Math.max(buf.byteLength * 2, n + value.byteLength)); b2.set(buf.subarray(0, n)); buf = b2; }
+    buf.set(value, n); n += value.byteLength; onBytes(value.byteLength, false);
+  }
+  buf = n === buf.byteLength ? buf : buf.slice(0, n);
+  if (cache) {
+    try { await cache.put(key, new Response(buf, { headers: { "content-type": "application/octet-stream", "content-length": String(n) } })); }
+    catch (e) { console.warn("cache put failed (quota?)", e); }
+  }
+  return { buf, cached: false };
+}
+
+export async function hasCachedModel(version = "001") {
+  try { const c = await caches.open(VERSIONS[version].cache); return (await c.keys()).length >= 5; } catch { return false; }
+}
+
+export async function clearCache() {
+  let ok = true;
+  for (const v of Object.values(VERSIONS)) { try { ok = (await caches.delete(v.cache)) && ok; } catch { ok = false; } }
+  return ok;
+}
+
+// The Hub counts a model download only on a request for the model repo's root config.json; the ONNX files
+// come from this Space, so every model load reads the loaded release's config.json once -- the same per-load
+// count a Python from_pretrained() produces, cached weights or not.
+function countLoad(version) { fetch(countUrl(version), { cache: "no-store" }).catch(() => {}); }
+
+// ---- the model ---------------------------------------------------------------------------------
+export class Agate {
+  // base: the page's models folder (001); override: an explicit models base for every version (local layout)
+  constructor({ base = "./models/", override = null, version = "003", ep = "webgpu", variant = null, optLevel = "all" } = {}) {
+    this.optLevel = optLevel;
+    const slash = (u) => (u.endsWith("/") ? u : u + "/");
+    this.version = version;
+    this.spec = VERSIONS[version];
+    this.base = override ? slash(override) + this.spec.dir : (this.spec.remote || slash(base) + this.spec.dir);
+    this.ep = ep;
+    this.variant = variant;          // e.g. "full16": fp16-compute generator (experimental, 001 only)
+    this.marks = marks(this.spec.release, REPO);
+  }
+
+  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 = () => {}) {
+    const meta = this.manifest.files[f];
+    return (await fetchCached(this.spec.cache, this.base + f, this.fileKey(f), meta.bytes, onBytes)).buf;
+  }
+
+  // onProgress({loaded, total, file, phase}); resolution: 003 only (512 default, or 256)
+  async load(onProgress = () => {}, resolution = null) {
+    const manifest = await (await fetch(this.base + "manifest.json", { cache: "no-cache" })).json();
+    this.manifest = manifest;
+    const swap = (this.variant && manifest.variants?.[this.variant]) || {};
+    const real = (f) => swap[f] || f;
+    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 || 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 = {};
+    const t0 = performance.now();
+    for (const f of need) {
+      const meta = manifest.files[real(f)];
+      const { buf } = await fetchCached(this.spec.cache, this.base + real(f), this.fileKey(real(f)), meta.bytes, (b, c) => {
+        loaded += b; if (c) fromCache += b; onProgress({ phase: "download", file: f, loaded, total });
+      });
+      bufs[f] = buf;
+    }
+    this.stats = { downloadMB: total / 1e6, fromCacheMB: fromCache / 1e6, downloadMs: performance.now() - t0 };
+    countLoad(this.version);
+    // drop cached files of older builds of this release (keys carry the sha256)
+    try {
+      const cache = await openCache(this.spec.cache), keep = new Set(need.map((f) => new URL(this.fileKey(real(f)).url, location.href).href));
+      if (cache) for (const req of await cache.keys()) if (!keep.has(req.url)) await cache.delete(req);
+    } catch (e) { console.warn("cache cleanup", e); }
+    const dec = new TextDecoder();
+    this.tok = new Tokenizer(JSON.parse(dec.decode(bufs["tokenizer.json"])), JSON.parse(dec.decode(bufs["tokenizer_config.json"])));
+
+    if (this.ep === "wasm") {
+      ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(8, navigator.hardwareConcurrency || 4) : 1;
+    }
+    const t1 = performance.now();
+    this.sessions = {};
+    for (const [name, f] of [["text", "text_encoder.onnx"], ["vae", "taesd_decoder.onnx"]]) {
+      onProgress({ phase: "init", file: f, loaded: total, total });
+      const ts = performance.now();
+      this.sessions[name] = await ort.InferenceSession.create(bufs[f], this.sessionOptions());
+      this.stats[`${name}SessionMs`] = performance.now() - ts;
+      bufs[f] = null;
+    }
+    onProgress({ phase: "init", file: "generator", loaded: total, total });
+    const ts = performance.now();
+    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 {
+      this.sessions.gen = await ort.InferenceSession.create(bufs["generator.onnx"], this.sessionOptions());
+    }
+    this.stats.genSessionMs = performance.now() - ts;
+    this.stats.sessionMs = performance.now() - t1;
+    return this.stats;
+  }
+
+  sessionOptions(extra = {}) { return { executionProviders: [this.ep], graphOptimizationLevel: this.optLevel, ...extra }; }
+
+  // 003: the generator graph for one resolution; the weights file is shared by both graphs.
+  async createGenerator(res, bufs = {}) {
+    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: ext, data: weights }], ...extra }));
+    this.resolution = Number(res);
+  }
+
+  async setResolution(res) {
+    if (!this.mr || Number(res) === this.resolution) return 0;
+    const t0 = performance.now();
+    await this.createGenerator(res);
+    return performance.now() - t0;
+  }
+
+  async release() {
+    for (const s of Object.values(this.sessions || {})) { try { await s?.release(); } catch { /* */ } }
+    this.sessions = null; this.weights = null;
+  }
+
+  tokenize(text) {
+    const max = this.manifest.text_max_len;
+    let ids = this.tok.encode(text).ids;                 // [CLS] ... [SEP], as the HF tokenizer
+    if (ids.length > max) ids = ids.slice(0, max - 1).concat(ids[ids.length - 1]);   // truncation=True
+    return ids;
+  }
+
+  async encode(text) {
+    const ids = this.tokenize(text);
+    const n = ids.length;
+    const feeds = {
+      input_ids: new ort.Tensor("int64", BigInt64Array.from(ids, BigInt), [1, n]),
+      attention_mask: new ort.Tensor("int64", new BigInt64Array(n).fill(1n), [1, n]),
+    };
+    const out = await this.sessions.text.run(feeds);
+    const h = out.last_hidden_state;
+    const data = h.data.slice();                          // (1, n, 512) fp32
+    h.dispose?.();
+    return { data, n, ids };
+  }
+
+  // The prompt as the model sees it: 003 runs its prompt pipeline (prompt.js), 001/002 use the prompt as typed.
+  prepare(prompt, negative = "") {
+    const pp = this.manifest.prompt_pipeline;
+    if (!pp) return { text: prompt, negative };
+    const [text, neg] = prepare(prompt, negative, pp.normalize, pp.spell);
+    return { text, negative: neg };
+  }
+
+  static gaussianNoise(seed, count) {
+    // mulberry32 + Box-Muller. Seeds do NOT reproduce PyTorch's torch.randn stream.
+    let a = seed >>> 0;
+    const rnd = () => { a |= 0; a = (a + 0x6D2B79F5) | 0; let t = Math.imul(a ^ (a >>> 15), 1 | a); t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t; return ((t ^ (t >>> 14)) >>> 0) / 4294967296; };
+    const out = new Float32Array(count);
+    for (let i = 0; i < count; i += 2) {
+      const u1 = Math.max(rnd(), 1e-12), u2 = rnd();
+      const r = Math.sqrt(-2 * Math.log(u1));
+      out[i] = r * Math.cos(2 * Math.PI * u2);
+      if (i + 1 < count) out[i + 1] = r * Math.sin(2 * Math.PI * u2);
+    }
+    return out;
+  }
+
+  // SD3's resolution shift in Agate's convention (t = 0 noise), as agate/pipeline.py shift_t
+  static shiftT(t, shift) {
+    if (shift === 1) return t;
+    const s = 1 - t;
+    return 1 - shift * s / (1 + (shift - 1) * s);
+  }
+
+  // -> { 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, 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);
+    const per = C * hw * hw;
+    const T = { };
+    let t0 = performance.now();
+    const prep = this.prepare(prompt, negative);
+    // cfg == 1 means v = v_cond: a batch of 1, no negative prompt (the distilled student's mode, as the patched
+    // agate/pipeline.py in the model repo does)
+    const single = cfg === 1 && !this.fast && !this.mr;
+    const nb = single ? 1 : 2;
+    const c = await this.encode(prep.text), u = single ? { data: new Float32Array(0), n: 0 } : await this.encode(prep.negative);
+    T.textMs = performance.now() - t0;
+    const BUCKETS = this.manifest.buckets;
+    const need = Math.max(c.n, u.n);
+    const L = BUCKETS.find((b) => b >= need) ?? BUCKETS[BUCKETS.length - 1];
+    const ctx = new Float32Array(nb * L * D), mask = new Float32Array(nb * L);
+    ctx.set(c.data, 0); mask.fill(1, 0, c.n);             // zero padding, as F.pad in the pipeline
+    if (!single) { ctx.set(u.data, L * D); mask.fill(1, L, L + u.n); }
+    const feeds = { ctx: new ort.Tensor("float32", ctx, [nb, L, D]), mask: new ort.Tensor("float32", mask, [nb, L]) };
+    let countsNonzero = 0;
+    if (this.mr) {                                        // count code: conditional half only
+      const cnt = new Float32Array(2 * L);
+      if (this.manifest.prompt_pipeline?.count_code) cnt.set(countVector(this.tok, prep.text, c.ids, Math.min(L, c.n)), 0);
+      countsNonzero = cnt.reduce((s, v) => s + (v > 0), 0);
+      feeds.counts = new ort.Tensor("float32", cnt, [2, L]);
+    }
+
+    let z = noise ? Float32Array.from(noise) : Agate.gaussianNoise(seed, per);
+    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(nb * per), tt = new Float32Array(nb);
+    for (let i = 0; i < steps; i++) {
+      if (shouldStop()) throw new Error("stopped");
+      const ts = performance.now();
+      zz.set(z, 0); if (!single) zz.set(z, per);
+      tt.fill(grid[i]);
+      const want = onPreview && this.hasPlan ? ["v", "plan"] : ["v"];
+      const out = await this.sessions.gen.run({ z: new ort.Tensor("float32", zz, [nb, C, hw, hw]), t: new ort.Tensor("float32", tt, [nb]), ...feeds }, want);
+      const v = out.v.data;
+      const dt = grid[i + 1] - grid[i];
+      const zn = new Float32Array(per);
+      let x1 = null;
+      if (out.plan) x1 = new Float32Array(per);
+      for (let k = 0; k < per; k++) {
+        const vc = v[k], vu = single ? vc : v[per + k], g = vu + cfg * (vc - vu);
+        zn[k] = z[k] + dt * g;
+        if (x1) x1[k] = z[k] + (1 - grid[i]) * g;
+      }
+      out.v.dispose?.();
+      if (out.plan) { const plan = out.plan.data.slice(); out.plan.dispose?.(); onPreview({ step: i + 1, steps, plan, x1, hw }); }
+      z = zn;
+      stepMs.push(performance.now() - ts);
+      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];
+
+    t0 = performance.now();
+    const img = await this.sessions.vae.run({ latent: new ort.Tensor("float32", z, [1, C, hw, hw]) });
+    const x = img.image.data, H = img.image.dims[2], W = img.image.dims[3];
+    const rgba = new Uint8ClampedArray(H * W * 4);
+    for (let p = 0; p < H * W; p++) {
+      for (let ch = 0; ch < 3; ch++) {
+        const val = Math.min(1, Math.max(-1, x[ch * H * W + p]));
+        rgba[p * 4 + ch] = Math.round((val + 1) * 127.5);
+      }
+      rgba[p * 4 + 3] = 255;
+    }
+    T.decodeMs = performance.now() - t0;
+    t0 = performance.now();
+    if (watermark) embedWatermark(rgba, W, H, this.marks.payload);
+    T.markMs = performance.now() - t0;
+    T.totalMs = T.textMs + T.samplerMs + T.decodeMs + T.markMs;
+    T.bucket = L;
+    return { rgba, width: W, height: H, latent: z, timings: T, prepared: { ...prep, countsNonzero }, watermarked: watermark };
+  }
+}

js/app.js

@@ -0,0 +1,169 @@
+import { Agate, VERSIONS, webgpuStatus, hasCachedModel } from "./agate.js";
+import { addPngText } from "./marking.js";
+
+const VERSION = "4step";
+const STEPS = 4, CFG = 1.0;                     // the student's sampler: 4 uniform Euler steps, guidance baked in
+const params = new URLSearchParams(location.search);
+let ep = params.get("ep") || "webgpu";
+
+const EXAMPLES = [
+  "a red cube on top of a blue sphere",
+  "a green teapot and a red cup on a table",
+  "a minimalist logo of a fox head, orange, flat design, white background",
+  "a cabin in a snowy forest at night, warm light in the windows",
+  "an astronaut riding a horse on the moon",
+  "a watercolor painting of a harbor with red sailboats",
+];
+
+const $ = (id) => document.getElementById(id);
+const ui = {
+  load: $("load"), loadrow: $("loadrow"), fill: $("bar-fill"), status: $("status"), statusR: $("status-r"),
+  prompt: $("prompt"), live: $("live"), seed: $("seed"), dice: $("dice"), go: $("go"),
+  canvas: $("canvas"), frame: $("frame"), placeholder: $("placeholder"), timings: $("timings"),
+  save: $("save"), aiNote: $("ai-note"),
+};
+const store = {
+  get(k) { try { return localStorage.getItem(k); } catch { return null; } },
+  set(k, v) { try { localStorage.setItem(k, v); } catch { /* storage blocked */ } },
+};
+
+let agate = null, ready = false, running = false, pending = false, count = 0;
+window.__agate = { state: "idle" };             // read by automated browser tests
+
+function setProgress(frac, left, right = "") {
+  ui.fill.style.width = `${Math.round(Math.max(0, Math.min(1, frac)) * 100)}%`;
+  ui.status.textContent = left;
+  ui.statusR.textContent = right;
+}
+
+function renderExamples() {
+  for (const p of EXAMPLES) {
+    const b = document.createElement("button");
+    b.className = "chip"; b.type = "button"; b.textContent = p;
+    b.onclick = () => { ui.prompt.value = p; request(); };
+    $("examples").appendChild(b);
+  }
+}
+
+function showTimings(T) {
+  const rows = [
+    ["image", `${Math.round(T.totalMs)} ms`, true],
+    ["text encoder", `${T.textMs.toFixed(0)} ms`],
+    [`sampler (${STEPS} steps)`, `${T.samplerMs.toFixed(0)} ms`],
+    ["per step", `${T.stepMs.toFixed(0)} ms`],
+    ["decode (TAESD)", `${T.decodeMs.toFixed(0)} ms`],
+    ["watermark", `${T.markMs.toFixed(0)} ms`],
+    ["backend", ep === "wasm" ? "CPU (WASM)" : "WebGPU"],
+    ["images this session", String(count)],
+  ];
+  ui.timings.replaceChildren(...rows.flatMap(([k, v, big]) => {
+    const dt = document.createElement("dt"), dd = document.createElement("dd");
+    dt.textContent = k; dd.textContent = v; if (big) dd.className = "big";
+    return [dt, dd];
+  }));
+}
+
+async function load() {
+  ui.load.disabled = true;
+  window.__agate.state = "loading";
+  const t0 = performance.now();
+  try {
+    agate = new Agate({ version: VERSION, ep });
+    await agate.load(({ phase, file, loaded, total }) => {
+      if (phase === "download") setProgress(loaded / total, `Downloading ${file}`, `${(loaded / 1e6).toFixed(0)} / ${(total / 1e6).toFixed(0)} MB`);
+      else setProgress(1, `Starting ${file}`);
+    });
+    // the first WebGPU run compiles shaders: pay it here, not on the first keystroke
+    setProgress(1, "Warming up the GPU");
+    await agate.generate({ prompt: "a photo", seed: 0, steps: STEPS, cfg: CFG, watermark: false });
+    ready = true;
+    ui.loadrow.hidden = true;
+    ui.go.disabled = false;
+    ui.placeholder.hidden = true;
+    const s = ((performance.now() - t0) / 1000).toFixed(1);
+    setProgress(0, `Ready (${ep === "wasm" ? "CPU" : "WebGPU"})`, `loaded in ${s} s`);
+    window.__agate = { state: "ready", loadMs: performance.now() - t0, ep };
+    request();
+  } catch (e) {
+    console.error(e);
+    ui.load.disabled = false;
+    setProgress(0, "Could not load the model", String(e.message || e));
+    window.__agate = { state: "error", error: String(e.message || e) };
+  }
+}
+
+// Coalescing runner: while one image is being made, keystrokes only mark the prompt as dirty; the newest prompt is
+// drawn as soon as the GPU is free, so the page never queues stale work.
+function request() {
+  if (!ready) return;
+  pending = true;
+  if (!running) run();
+}
+
+async function run() {
+  running = true;
+  ui.frame.classList.add("busy");
+  while (pending) {
+    pending = false;
+    const prompt = ui.prompt.value.trim();
+    if (!prompt) continue;
+    const seed = Math.max(0, Math.min(4294967295, Number(ui.seed.value) || 0));
+    try {
+      window.__agate.state = "generating";
+      const r = await agate.generate({ prompt, seed, steps: STEPS, cfg: CFG, watermark: true });
+      const ctx = ui.canvas.getContext("2d");
+      ui.canvas.width = r.width; ui.canvas.height = r.height;
+      ctx.putImageData(new ImageData(r.rgba, r.width, r.height), 0, 0);
+      count += 1;
+      showTimings(r.timings);
+      ui.save.disabled = false;
+      ui.aiNote.hidden = false;
+      window.__agate = { state: "done", timings: r.timings, count, prompt };
+    } catch (e) {
+      console.error(e);
+      setProgress(0, "Error", String(e.message || e));
+      window.__agate = { state: "error", error: String(e.message || e) };
+    }
+  }
+  ui.frame.classList.remove("busy");
+  running = false;
+}
+
+async function save() {
+  const blob = await new Promise((r) => ui.canvas.toBlob(r, "image/png"));
+  const bytes = addPngText(new Uint8Array(await blob.arrayBuffer()), agate.marks.info);
+  const a = document.createElement("a");
+  a.href = URL.createObjectURL(new Blob([bytes], { type: "image/png" }));
+  a.download = `agate-4step-${ui.seed.value}.png`;
+  a.click();
+  setTimeout(() => URL.revokeObjectURL(a.href), 1000);
+}
+
+async function init() {
+  renderExamples();
+  ui.live.checked = store.get("agate4-live") !== "0";
+  ui.live.onchange = () => store.set("agate4-live", ui.live.checked ? "1" : "0");
+  ui.prompt.addEventListener("input", () => { if (ui.live.checked) request(); });
+  ui.prompt.addEventListener("keydown", (e) => { if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) request(); });
+  ui.seed.addEventListener("input", () => { if (ui.live.checked) request(); });
+  ui.dice.onclick = () => { ui.seed.value = String(Math.floor(Math.random() * 1e9)); request(); };
+  ui.go.onclick = request;
+  ui.load.onclick = load;
+  ui.save.onclick = save;
+
+  if (ep !== "wasm") {
+    const s = await webgpuStatus();
+    if (!s.ok) {
+      $("nogpu").hidden = false;
+      $("nogpu-why").textContent = s.why;
+      ui.load.disabled = true;
+      $("use-wasm").onclick = () => { ep = "wasm"; $("nogpu").hidden = true; ui.load.disabled = false; ui.load.textContent = "Load model on the CPU (about 615 MB)"; };
+      window.__agate = { state: "no-webgpu", why: s.why };
+      return;
+    }
+    ui.statusR.textContent = s.adapter || "";
+  }
+  if (params.has("autoload") || (await hasCachedModel(VERSION))) load();
+}
+
+init();

js/marking.js

@@ -0,0 +1,188 @@
+// AI-generated content marking for the images this page makes (EU AI Act Art. 50(2)), the same marks as the
+// Python packages (agate/marking.py of Logolabs/agate-preview-002 / -003):
+//  * an invisible watermark in the pixels: invisible-watermark's (MIT) 'dwtDctSvd' method, fixed 64-bit payload
+//    "AGATE" + release ("AGATE001" / "AGATE002" / "AGATE003"), ported from the library with its exact conventions
+//    (OpenCV 8-bit YUV, U channel only, Haar LL band, 4x4 DCT + SVD, s0 quantised to 36, the library's swapped
+//    detail order in the inverse DWT, numpy's truncating uint8 cast). agate.detect_watermark() in Python reads it.
+//  * provenance text chunks (tEXt) in the downloaded PNG: ai_generated, generator, model, watermark -- the keys
+//    and values the Python packages write. The prompt is NOT written.
+// Neither mark is tamper-proof; see the model cards.
+
+export const METHOD = "dwtDctSvd";
+const SCALE = 36, BLOCK = 4;
+const S2 = 0.7071067811865476;                      // PyWavelets' Haar tap, 1/sqrt(2)
+
+// payload keeps the base release (AGATE002), as agate/marking.py in ysharma/agate-002-4step does; the provenance
+// fields name the distilled model
+export function marks(release, repo = `Logolabs/agate-preview-${release}`) {
+  const payload = "AGATE" + release;               // 8 ASCII bytes = 64 bits
+  return {
+    payload,
+    info: {
+      ai_generated: "true",
+      generator: `Agate 002 4-step, unofficial distillation of Agate Preview ${release} (LogoLabs)`,
+      model: repo,
+      watermark: `invisible-watermark ${METHOD}, payload ${payload}`,
+    },
+  };
+}
+
+const desc = (x) => (x + 8192) >> 14;               // OpenCV CV_DESCALE(x, 14)
+const sat = (x) => (x < 0 ? 0 : x > 255 ? 255 : x);
+const u8unsafe = (x) => { const t = Math.trunc(x) % 256; return t < 0 ? t + 256 : t; };   // numpy float -> uint8
+
+// orthonormal 4-point DCT-II matrix (cv2.dct on a 4x4 block = C X C^T)
+const C = [];
+for (let k = 0; k < BLOCK; k++) {
+  C.push([]);
+  for (let n = 0; n < BLOCK; n++) C[k].push(Math.sqrt(k === 0 ? 1 / BLOCK : 2 / BLOCK) * Math.cos(Math.PI * (2 * n + 1) * k / (2 * BLOCK)));
+}
+function mul(A, B) { const R = [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]; for (let i = 0; i < 4; i++) for (let j = 0; j < 4; j++) { let s = 0; for (let k = 0; k < 4; k++) s += A[i][k] * B[k][j]; R[i][j] = s; } return R; }
+const T = (A) => A[0].map((_, j) => A.map((r) => r[j]));
+const CT = T(C);
+
+// largest singular value of a 4x4 matrix and its singular vectors (Jacobi on D^T D)
+function topSVD(D) {
+  const M = mul(T(D), D);
+  const V = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]];
+  for (let sweep = 0; sweep < 30; sweep++) {
+    let off = 0;
+    for (let p = 0; p < 4; p++) for (let q = p + 1; q < 4; q++) off += M[p][q] * M[p][q];
+    if (off < 1e-30) break;
+    for (let p = 0; p < 4; p++) for (let q = p + 1; q < 4; q++) {
+      if (Math.abs(M[p][q]) < 1e-300) continue;
+      const th = (M[q][q] - M[p][p]) / (2 * M[p][q]);
+      const t = Math.sign(th || 1) / (Math.abs(th) + Math.sqrt(th * th + 1));
+      const c = 1 / Math.sqrt(t * t + 1), s = t * c;
+      for (let k = 0; k < 4; k++) { const a = M[k][p], b = M[k][q]; M[k][p] = c * a - s * b; M[k][q] = s * a + c * b; }
+      for (let k = 0; k < 4; k++) { const a = M[p][k], b = M[q][k]; M[p][k] = c * a - s * b; M[q][k] = s * a + c * b; }
+      for (let k = 0; k < 4; k++) { const a = V[k][p], b = V[k][q]; V[k][p] = c * a - s * b; V[k][q] = s * a + c * b; }
+    }
+  }
+  let best = 0;
+  for (let i = 1; i < 4; i++) if (M[i][i] > M[best][best]) best = i;
+  const s0 = Math.sqrt(Math.max(0, M[best][best]));
+  let v = V.map((r) => r[best]), u;
+  if (s0 < 1e-12) { u = [1, 0, 0, 0]; v = [1, 0, 0, 0]; }
+  else u = D.map((r) => (r[0] * v[0] + r[1] * v[1] + r[2] * v[2] + r[3] * v[3]) / s0);
+  return { s0, u, v };
+}
+
+function payloadBits(payload) {
+  const bytes = new TextEncoder().encode(payload), bits = [];
+  for (const b of bytes) for (let k = 7; k >= 0; k--) bits.push((b >> k) & 1);
+  return bits;
+}
+
+// U channel (OpenCV BGR2YUV, 8 bit) of an RGBA buffer -> Int32Array(W*H), plus Y and V for the inverse
+function toYUV(rgba, W, H) {
+  const Y = new Int32Array(W * H), U = new Int32Array(W * H), V = new Int32Array(W * H);
+  for (let p = 0; p < W * H; p++) {
+    const r = rgba[4 * p], g = rgba[4 * p + 1], b = rgba[4 * p + 2];
+    const y = desc(b * 1868 + g * 9617 + r * 4899);
+    Y[p] = y; U[p] = sat(desc((b - y) * 8061 + (128 << 14))); V[p] = sat(desc((r - y) * 14369 + (128 << 14)));
+  }
+  return { Y, U, V };
+}
+
+// LL band of the Haar DWT of U over the (H//4*4, W//4*4) region; -> {ca, ch, cv, cd} as Float64Array (h2 x w2)
+function haar(U, W, R, Cc) {
+  const h2 = R / 2, w2 = Cc / 2, n = h2 * w2;
+  const ca = new Float64Array(n), ch = new Float64Array(n), cv = new Float64Array(n), cd = new Float64Array(n);
+  for (let i = 0; i < h2; i++) for (let j = 0; j < w2; j++) {
+    const a = U[(2 * i) * W + 2 * j], b = U[(2 * i) * W + 2 * j + 1], c = U[(2 * i + 1) * W + 2 * j], d = U[(2 * i + 1) * W + 2 * j + 1];
+    const k = i * w2 + j;
+    // PyWavelets' exact float order: pairs along rows first (a|c, b|d), then along columns
+    const l0 = a * S2 + c * S2, l1 = b * S2 + d * S2, h0 = a * S2 - c * S2, h1 = b * S2 - d * S2;
+    ca[k] = l0 * S2 + l1 * S2; cv[k] = l0 * S2 - l1 * S2; ch[k] = h0 * S2 + h1 * S2; cd[k] = h0 * S2 - h1 * S2;
+  }
+  return { ca, ch, cv, cd, h2, w2 };
+}
+
+function blockAt(ca, w2, bi, bj) { const B = []; for (let r = 0; r < 4; r++) { B.push([]); for (let c = 0; c < 4; c++) B[r].push(ca[(bi * 4 + r) * w2 + bj * 4 + c]); } return B; }
+
+// Embed `payload` in place into rgba (Uint8ClampedArray / Uint8Array, W*H*4). Needs W*H >= 256*256.
+export function embedWatermark(rgba, W, H, payload) {
+  if (W * H < 256 * 256) throw new Error("watermark: image too small (needs at least 256 x 256)");
+  const bits = payloadBits(payload);
+  const { Y, U, V } = toYUV(rgba, W, H);
+  const R = Math.floor(H / 4) * 4, Cc = Math.floor(W / 4) * 4;
+  const { ca, ch, cv, cd, h2, w2 } = haar(U, W, R, Cc);
+  const nbi = Math.floor(h2 / 4), nbj = Math.floor(w2 / 4);
+  let num = 0;
+  for (let bi = 0; bi < nbi; bi++) for (let bj = 0; bj < nbj; bj++, num++) {
+    const D = mul(mul(C, blockAt(ca, w2, bi, bj)), CT);
+    const { s0, u, v } = topSVD(D);
+    const s1 = (Math.floor(s0 / SCALE) + 0.25 + 0.5 * bits[num % bits.length]) * SCALE;
+    const Dn = D.map((row, i) => row.map((x, j) => x + (s1 - s0) * u[i] * v[j]));
+    const Bn = mul(mul(CT, Dn), C);
+    for (let r = 0; r < 4; r++) for (let c = 0; c < 4; c++) ca[(bi * 4 + r) * w2 + bj * 4 + c] = Bn[r][c];
+  }
+  // inverse Haar with the library's swapped details (cv as H, ch as V), truncating uint8 cast
+  for (let i = 0; i < h2; i++) for (let j = 0; j < w2; j++) {
+    // pywt.idwt2((ca, (cv, ch, cd))): the library's swapped details; PyWavelets' float order (columns, then rows)
+    const k = i * w2 + j, A = ca[k], hIn = cv[k], vIn = ch[k], d = cd[k];
+    const L0 = A * S2 + vIn * S2, L1 = A * S2 - vIn * S2, H0 = hIn * S2 + d * S2, H1 = hIn * S2 - d * S2;
+    U[(2 * i) * W + 2 * j] = u8unsafe(L0 * S2 + H0 * S2);
+    U[(2 * i) * W + 2 * j + 1] = u8unsafe(L1 * S2 + H1 * S2);
+    U[(2 * i + 1) * W + 2 * j] = u8unsafe(L0 * S2 - H0 * S2);
+    U[(2 * i + 1) * W + 2 * j + 1] = u8unsafe(L1 * S2 - H1 * S2);
+  }
+  for (let p = 0; p < W * H; p++) {                  // OpenCV YUV2BGR, 8 bit
+    const y = Y[p], uu = U[p] - 128, vq = V[p] - 128;
+    rgba[4 * p + 2] = sat(y + desc(uu * 33292));
+    rgba[4 * p + 1] = sat(y + desc(uu * -6472 + vq * -9519));
+    rgba[4 * p] = sat(y + desc(vq * 18678));
+  }
+  return rgba;
+}
+
+// -> {bits, text, bitAccuracy} against `payload` (a self-check; the reference detector is the Python one)
+export function readWatermark(rgba, W, H, payload) {
+  const want = payloadBits(payload), n = want.length;
+  const { U } = toYUV(rgba, W, H);
+  const R = Math.floor(H / 4) * 4, Cc = Math.floor(W / 4) * 4;
+  const { ca, h2, w2 } = haar(U, W, R, Cc);
+  const nbi = Math.floor(h2 / 4), nbj = Math.floor(w2 / 4);
+  const sum = new Float64Array(n), cnt = new Float64Array(n);
+  let num = 0;
+  for (let bi = 0; bi < nbi; bi++) for (let bj = 0; bj < nbj; bj++, num++) {
+    const { s0 } = topSVD(mul(mul(C, blockAt(ca, w2, bi, bj)), CT));
+    sum[num % n] += (s0 % SCALE) > SCALE * 0.5 ? 1 : 0; cnt[num % n] += 1;
+  }
+  const bits = Array.from(sum, (s, k) => (s / cnt[k]) * 255 > 127 ? 1 : 0);
+  let same = 0; bits.forEach((b, k) => { if (b === want[k]) same++; });
+  const bytes = []; for (let k = 0; k < n; k += 8) { let b = 0; for (let q = 0; q < 8; q++) b = (b << 1) | bits[k + q]; bytes.push(b); }
+  return { bits, text: String.fromCharCode(...bytes), bitAccuracy: same / n };
+}
+
+// ---- PNG tEXt chunks ---------------------------------------------------------------------------
+let CRC_TABLE = null;
+function crc32(bytes) {
+  if (!CRC_TABLE) { CRC_TABLE = new Uint32Array(256); for (let n = 0; n < 256; n++) { let c = n; for (let k = 0; k < 8; k++) c = c & 1 ? 0xedb88320 ^ (c >>> 1) : c >>> 1; CRC_TABLE[n] = c >>> 0; } }
+  let c = 0xffffffff;
+  for (const b of bytes) c = CRC_TABLE[(c ^ b) & 0xff] ^ (c >>> 8);
+  return (c ^ 0xffffffff) >>> 0;
+}
+
+// PNG bytes -> PNG bytes with one tEXt chunk per entry of `info` (Latin-1 keys/values), inserted before IEND.
+export function addPngText(png, info) {
+  const latin1 = (s) => Uint8Array.from(s, (ch) => { const c = ch.charCodeAt(0); return c < 256 ? c : 63; });
+  const chunks = [];
+  for (const [k, v] of Object.entries(info)) {
+    const data = new Uint8Array([...latin1(k), 0, ...latin1(String(v))]);
+    const typeData = new Uint8Array([116, 69, 88, 116, ...data]);          // "tEXt"
+    const out = new Uint8Array(12 + data.length), dv = new DataView(out.buffer);
+    dv.setUint32(0, data.length); out.set(typeData, 4); dv.setUint32(8 + data.length, crc32(typeData));
+    chunks.push(out);
+  }
+  // find IEND (the last chunk): 12 bytes from the end in every well-formed PNG
+  const iend = png.length - 12;
+  const extra = chunks.reduce((s, c) => s + c.length, 0);
+  const res = new Uint8Array(png.length + extra);
+  res.set(png.subarray(0, iend), 0);
+  let o = iend;
+  for (const c of chunks) { res.set(c, o); o += c.length; }
+  res.set(png.subarray(iend), o);
+  return res;
+}

js/prompt.js

@@ -0,0 +1,249 @@
+// Agate Preview 003 prompt pipeline, a line-by-line port of agate/prompt_norm.py (+ AgatePipeline.prepare
+// and prompt_norm.count_tensor) of Logolabs/agate-preview-003. Checked against the Python package by
+// test/prompt_golden.json (tools: export/prompt_golden.py writes it, test/prompt_golden.mjs or ?golden=1 checks it).
+//
+// Python-compatibility notes (the reason for the helpers below):
+//  * Python's \s, \d, \b and str.isalnum/isspace are Unicode-aware; JS's are ASCII. PY_WS, \p{Nd} and the
+//    Unicode word-boundary WB mirror Python's definitions (\w = letters, numbers, underscore).
+//  * Python counts code points; JS strings count UTF-16 units. Lengths that gate behaviour use cpLen().
+//  * Token offsets come from the HF fast tokenizer in Python; here they are rebuilt from the byte-level BPE
+//    tokens (tokenOffsets), in UTF-16 units, which gives the same overlaps as Python's code-point offsets.
+
+const PY_WS = "\\t\\n\\x0b\\x0c\\r\\x1c-\\x20\\x85\\xa0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000";
+const S = `[${PY_WS}]`;                         // Python's \s
+const NS = `[^${PY_WS}]`;                       // Python's \S
+const W = "[\\p{L}\\p{N}_]";                    // Python's \w (str patterns)
+const WB = `(?:(?<=${W})(?!${W})|(?<!${W})(?=${W}))`;   // Python's \b
+const D = "\\p{Nd}";                            // Python's \d
+const re = (src, flags = "") => new RegExp(src, flags.includes("u") ? flags : flags + "u");
+
+const NUM_LIST = "zero one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen nineteen twenty".split(" ");
+const NUM_WORDS = new Map(NUM_LIST.map((w, i) => [w, i]));
+NUM_WORDS.set("dozen", 12);
+const WORD_OF = new Map(NUM_LIST.map((w, i) => [i, w]));
+const NOT_COUNT_NEXT = new Set(["pm", "am", "o'clock", "percent", "%", "years", "year", "hours", "hour", "minutes", "minute",
+  "seconds", "second", "times", "px", "cm", "mm", "m", "km", "kg", "g", "inch", "inches", "feet",
+  "degrees", "x", "d", "k", "th", "st", "nd", "rd", "of", "o"]);
+const ONE_PRONOUN_PREV = new Set(["the", "this", "that", "no", "each", "every", "any", "some", "which", "another", "someone",
+  "everyone", "anyone", "only"]);
+export const SPELL_SEP = " || spell: ";
+const SPELL_MAX_FRAC = 0.6, SPELL_MAX_CHARS = 40, SPELL_MAX_WORDS = 6, SPELL_MAX_SPANS = 3;
+
+const QUOTE_SRC = '"[^"]*"|“[^”]*”';
+const cpLen = (s) => { let n = 0; for (const _ of s) n++; return n; };     // len() in Python
+const isSpace = (ch) => re(`^${S}$`).test(ch);
+const pyStrip = (s, chars = null) => {
+  const f = chars === null ? isSpace : (c) => chars.includes(c);
+  let a = 0, b = s.length;
+  while (a < b && f(s[a])) a++;
+  while (b > a && f(s[b - 1])) b--;
+  return s.slice(a, b);
+};
+const pySplit = (s) => s.split(re(`${S}+`)).filter((x) => x.length > 0);   // str.split()
+const isAlnum = (ch) => /^[\p{L}\p{N}]$/u.test(ch);
+const isAlpha = (w) => w.length > 0 && /^\p{L}+$/u.test(w);
+const isUpper = (w) => w !== w.toLowerCase() && w === w.toUpperCase();     // ASCII words only reach these
+const isLower = (w) => w !== w.toUpperCase() && w === w.toLowerCase();
+const isDigits = (w) => /^\p{Nd}+$/u.test(w);
+
+function digitValue(cp) {              // value of one \p{Nd} code point: Nd code points come in runs of ten
+  let start = cp;
+  while (/\p{Nd}/u.test(String.fromCodePoint(start - 1))) start--;
+  return (cp - start) % 10;
+}
+function pyInt(w) {                    // int() of a \p{Nd}+ string
+  let n = 0;
+  for (const ch of w) n = n * 10 + digitValue(ch.codePointAt(0));
+  return n;
+}
+
+function outsideQuotes(text) {
+  const spans = [];
+  let last = 0;
+  for (const m of text.matchAll(re(QUOTE_SRC, "g"))) { spans.push([last, m.index]); last = m.index + m[0].length; }
+  spans.push([last, text.length]);
+  return spans.filter(([a, b]) => b > a);
+}
+
+function mapOutside(text, fn) {
+  let out = "", last = 0;
+  for (const [a, b] of outsideQuotes(text)) { out += text.slice(last, a) + fn(text.slice(a, b)); last = b; }
+  return out + text.slice(last);
+}
+
+function fixCase(words) {
+  const long2 = words.filter((w) => cpLen(w) >= 2 && isAlpha(w));
+  if (long2.length && long2.filter(isUpper).length / long2.length >= 0.6) return "shout";
+  const long3 = words.filter((w) => cpLen(w) >= 3 && isAlpha(w));
+  if (long3.length >= 3) {
+    const t = long3.filter((w) => { const c = [...w]; return isUpper(c[0]) && isLower(c.slice(1).join("")); }).length;
+    if (t / long3.length >= 0.7) return "title";
+  }
+  return null;
+}
+
+export function isJsonCaption(text) {
+  const s = pyStrip(text.split(SPELL_SEP)[0]);
+  if (!(s.startsWith("{") && s.endsWith("}"))) return false;
+  try { const v = JSON.parse(s); return v !== null && typeof v === "object" && !Array.isArray(v); } catch { return false; }
+}
+
+function isCountContext(seg, end) {
+  const nxt = re(`^${S}*([A-Za-z%']+)`).exec(seg.slice(end));
+  return !!nxt && !NOT_COUNT_NEXT.has(nxt[1].toLowerCase());
+}
+
+function canonNumbers(seg) {
+  return seg.replace(re(`${WB}[A-Za-z]+${WB}|${WB}${D}+${WB}`, "g"), (w, off) => {
+    const lw = w.toLowerCase();
+    if (NUM_WORDS.has(lw)) return lw;
+    if (isDigits(w)) { const n = pyInt(w); if (n >= 1 && n <= 20 && isCountContext(seg, off + w.length)) return WORD_OF.get(n); }
+    return w;
+  });
+}
+
+export function normalize(text) {
+  if (isJsonCaption(text)) return text;
+  text = pyStrip(text.replace(re(`${S}+`, "g"), " "));
+  const words = [];
+  for (const [a, b] of outsideQuotes(text)) for (const m of text.slice(a, b).matchAll(re(`[A-Za-z][A-Za-z'\\-]*|${D}+`, "g"))) words.push(m[0]);
+  const mode = fixCase(words);
+  if (mode === "shout") text = mapOutside(text, (s) => s.toLowerCase());
+  else if (mode === "title") {
+    const first = re(`^${S}*${NS}+`).exec(text);
+    const head = first ? first[0].length : 0;
+    text = text.slice(0, head) + mapOutside(text.slice(head), (s) => s.replace(re(`${WB}[A-Z][a-z]*${WB}`, "g"), (m) => m.toLowerCase()));
+  }
+  return mapOutside(text, canonNumbers);
+}
+
+export function findCounts(text) {
+  const out = [];
+  if (isJsonCaption(text)) return out;
+  text = text.split(SPELL_SEP)[0];
+  for (const [a, b] of outsideQuotes(text)) {
+    const seg = text.slice(a, b);
+    const pat = re(`${WB}(?:a${S}+)?(dozen|pair(?=${S}+of${WB}))${WB}|${WB}([A-Za-z]+|${D}+)${WB}`, "gd");
+    for (const m of seg.matchAll(pat)) {
+      if (m[1] !== undefined) {
+        const [s1, e1] = m.indices[1];
+        out.push([a + s1, a + e1, m[1].toLowerCase() === "dozen" ? 12 : 2]);
+        continue;
+      }
+      const w = m[2], lw = w.toLowerCase();
+      const [s2, e2] = m.indices[2];
+      let n;
+      if (isDigits(lw)) { n = pyInt(lw); if (!(n >= 1 && n <= 99)) continue; }
+      else if (NUM_WORDS.has(lw) && lw !== "zero" && lw !== "dozen") n = NUM_WORDS.get(lw);
+      else continue;
+      if (!isCountContext(seg, m.index + m[0].length)) continue;
+      const prev = seg.slice(0, m.index).match(/[A-Za-z]+/g) || [];
+      if (lw === "one" && prev.length && ONE_PRONOUN_PREV.has(prev[prev.length - 1].toLowerCase())) continue;
+      out.push([a + s2, a + e2, n]);
+    }
+  }
+  return out;
+}
+
+export function textSpans(text) {
+  const body = text.split(SPELL_SEP)[0];
+  const out = [];
+  for (const m of body.matchAll(re(QUOTE_SRC, "g"))) {
+    const c = pyStrip(m[0].slice(1, -1));
+    if (c && [...c].some(isAlnum) && cpLen(c) <= SPELL_MAX_FRAC * cpLen(body) && cpLen(c) <= SPELL_MAX_CHARS && pySplit(c).length <= SPELL_MAX_WORDS) out.push(c);
+  }
+  return out.slice(0, SPELL_MAX_SPANS);
+}
+
+export function spell(content) {
+  const words = pySplit(content).map((w) => [...w].filter((ch) => isAlnum(ch) || "&!?'-".includes(ch)));
+  return words.filter((w) => w.length).map((w) => w.join(" ")).join(" / ");
+}
+
+export function addSpelling(text) {
+  if (text.includes(SPELL_SEP) || isJsonCaption(text)) return text;
+  const spans = textSpans(text);
+  if (!spans.length) return text;
+  return text + SPELL_SEP + spans.map(spell).join(" ; ");
+}
+
+export function splitNegatives(text) {
+  const negs = [];
+  if (isJsonCaption(text)) return [text, negs];
+  const pat = re(`${S}*,?${S}*${WB}(?:with${S}+)?(?:without|no)${S}+(?:any${S}+)?(?:a${S}+|an${S}+|the${S}+)?([^,.;"“]+?)(?=${S}+and${S}+|[,.;]|$)`, "gi");
+  const cut = (seg) => seg.replace(pat, (...g) => { negs.push(pyStrip(g[1])); return ""; });
+  let pos = mapOutside(text, cut);
+  pos = pyStrip(pos.replace(re(`${S}+([,.;])`, "g"), "$1").replace(re(`${S}+`, "g"), " "), " ,;");
+  return [pos, negs];
+}
+
+// AgatePipeline.prepare: -> [prompt as the text encoder sees it, negative prompt]
+export function prepare(prompt, negativePrompt = "", normalizeOn = true, spellOn = true) {
+  let negs = [];
+  if (normalizeOn) [prompt, negs] = splitNegatives(normalize(prompt));
+  if (spellOn) prompt = addSpelling(prompt);
+  const negative = [pyStrip(negativePrompt), ...negs].filter((n) => n).join(", ");
+  return [prompt, negative];
+}
+
+// ---- token offsets for a byte-level BPE tokenizer (tokenizers.js has no offset mapping) -----------
+let BYTE_OF = null;                       // GPT-2 bytes_to_unicode, inverted: char -> byte
+function byteDecoder() {
+  if (BYTE_OF) return BYTE_OF;
+  const bs = [];
+  for (let b = 33; b <= 126; b++) bs.push(b);
+  for (let b = 161; b <= 172; b++) bs.push(b);
+  for (let b = 174; b <= 255; b++) bs.push(b);
+  const cs = bs.slice();
+  let n = 0;
+  for (let b = 0; b < 256; b++) if (!bs.includes(b)) { bs.push(b); cs.push(256 + n); n++; }
+  BYTE_OF = new Map(bs.map((b, i) => [String.fromCodePoint(cs[i]), b]));
+  return BYTE_OF;
+}
+
+// ids: the token ids of `text` ([CLS] ... [SEP]). -> [[start, end], ...] in UTF-16 units of text, [0, 0] for
+// special tokens: Python's offset_mapping (code points) mapped to JS string indices.
+export function tokenOffsets(tok, text, ids) {
+  const dec = byteDecoder();
+  const nfc = text.normalize("NFC");
+  const enc = new TextEncoder();
+  // byte index -> [UTF-16 start, UTF-16 end] of the character that byte belongs to
+  const charOfByte = [];
+  let u = 0;
+  for (const ch of nfc) { const nb = enc.encode(ch).length; for (let k = 0; k < nb; k++) charOfByte.push([u, u + ch.length]); u += ch.length; }
+  const added = tok.get_added_tokens_decoder?.() ?? new Map();
+  const out = [];
+  let pos = 0;
+  for (const id of ids) {
+    const at = added.get(Number(id));
+    let nbytes;
+    if (at && at.special) { out.push([0, 0]); continue; }
+    if (at) nbytes = enc.encode(at.content).length;
+    else {
+      const s = tok.id_to_token(Number(id));
+      nbytes = 0;
+      for (const ch of s) nbytes += dec.has(ch) ? 1 : enc.encode(ch).length;
+    }
+    if (nbytes === 0 || pos >= charOfByte.length) { out.push([0, 0]); continue; }
+    const a = charOfByte[pos][0], b = charOfByte[Math.min(pos + nbytes, charOfByte.length) - 1][1];
+    out.push([a, b]);
+    pos += nbytes;
+  }
+  // text that NFC changed: offsets refer to the NFC string; the count spans are found in the NFC string too
+  return out;
+}
+
+// prompt_norm.count_tensor for one prompt: Float32Array(length), the count value at the tokens that spell one.
+export function countVector(tok, text, ids, length) {
+  const out = new Float32Array(length);
+  const nfc = text.normalize("NFC");
+  const spans = findCounts(nfc);
+  if (!spans.length) return out;
+  const offs = tokenOffsets(tok, nfc, ids);
+  for (let j = 0; j < Math.min(length, offs.length); j++) {
+    const [a, b] = offs[j];
+    if (b <= a) continue;
+    for (const [s, e, n] of spans) if (a < e && b > s) out[j] = n;
+  }
+  return out;
+}