app.js

// app.js — the ultimate input field on WebGPU, with bankML beside it as a separate engine.
//   generic mode (GPU): Agate draws (onnxruntime-web on WebGPU, Logolabs' MIT code and models) or Qwen3-0.6B writes
//                       (transformers.js on WebGPU) — the ⧉ GPU button chooses;
//   bankML mode (CPU):  bankML's own engine compiled to WebAssembly (browseML), sha256-verified, a receipt on every answer.
// Nothing leaves the browser except the model downloads, which are cached after the first time.
import * as browseml from "./browseml.js";

const $ = (id) => document.getElementById(id);
const secs = (s) => (s < 60 ? `${s.toFixed(1)} s` : `${Math.floor(s / 60)} min ${Math.round(s % 60)} s`);
const mb = (b) => `${(b / 1e6).toFixed(0)} MB`;
const sha256hex = async (data) => [...new Uint8Array(await crypto.subtle.digest("SHA-256", typeof data === "string" ? new TextEncoder().encode(data) : data))]
  .map((b) => b.toString(16).padStart(2, "0")).join("");

// ── settings, kept in this browser ───────────────────────────────────────────────────────────────────────────
const KEY = "ultimate-webgpu:settings";
const DEFAULTS = { gpu: "qwen", tokens: 384, temp: 0.7, steps: 50, cfg: 3.0, res: 512, seed: null };
let settings = { ...DEFAULTS };
try { settings = { ...DEFAULTS, ...JSON.parse(localStorage.getItem(KEY) || "{}") }; } catch { /* storage off */ }
const save = () => { try { localStorage.setItem(KEY, JSON.stringify(settings)); } catch { /* storage off */ } showEngines(); };

// ── WebGPU on this machine ───────────────────────────────────────────────────────────────────────────────────
let gpu = { ok: false, why: "not checked" };
async function checkGpu() {
  if (!("gpu" in navigator)) gpu = { ok: false, why: "this browser has no WebGPU (navigator.gpu is missing): the GPU engines fall back to the CPU, slowly" };
  else {
    try {
      const a = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
      if (!a) gpu = { ok: false, why: "WebGPU is present but no GPU adapter is available" };
      else {
        const info = a.info || {};
        gpu = { ok: true, name: [info.vendor, info.architecture, info.description].filter(Boolean).join(" ") || "a GPU", f16: a.features.has("shader-f16"), maxBuffer: a.limits.maxBufferSize };
      }
    } catch (e) { gpu = { ok: false, why: `WebGPU adapter request failed: ${e.message || e}` }; }
  }
  $("gpu").textContent = gpu.ok ? `WebGPU: ${gpu.name}${gpu.f16 ? " · shader-f16" : ""}` : `No WebGPU: ${gpu.why}`;
  $("gpu").className = `gpu ${gpu.ok ? "ok" : "no"}`;
}

// ── engine 1, GPU: Agate (text → image), loaded from Logolabs' repos and cached ─────────────────────────────────
let agateMod = null, agate = null, agateLoading = null;
async function loadAgate(progress) {
  if (agate) return agate;
  if (agateLoading) return agateLoading;
  agateLoading = (async () => {
    agateMod = agateMod || (await import("./agate/agate.js"));
    const a = new agateMod.Agate({ version: "003", ep: gpu.ok ? "webgpu" : "wasm" });
    const t0 = performance.now();
    await a.load(({ phase, file, loaded, total }) => {
      progress(phase === "download" ? `Agate: downloading ${file}, ${mb(loaded)} of ${mb(total)} (once; then from this browser's cache)` : `Agate: preparing ${file} on ${gpu.ok ? "WebGPU" : "the CPU"}…`);
    }, settings.res);
    a.loadMs = performance.now() - t0;
    agate = a;
    return a;
  })().finally(() => { agateLoading = null; });
  return agateLoading;
}
async function askAgate(prompt, progress) {
  const a = await loadAgate(progress);
  const seed = settings.seed ?? Math.floor(Math.random() * 2 ** 31);
  const t0 = performance.now();
  const r = await a.generate({ prompt, seed, steps: settings.steps, cfg: settings.cfg, watermark: true,
    onStep: (i, n) => progress(`Agate on ${gpu.ok ? "WebGPU" : "the CPU"}: step ${i} of ${n}`) });
  const c = document.createElement("canvas"); c.width = r.width; c.height = r.height;
  c.getContext("2d").putImageData(new ImageData(r.rgba, r.width, r.height), 0, 0);
  const png = await new Promise((ok) => c.toBlob(ok, "image/png"));
  const bytes = agateMod ? (await import("./agate/marking.js")).addPngText(new Uint8Array(await png.arrayBuffer()), a.marks.info) : new Uint8Array(await png.arrayBuffer());
  const s = (performance.now() - t0) / 1000;
  return { bytes, width: r.width, height: r.height, seed, seconds: s, timings: r.timings, prepared: r.prepared };
}

// ── engine 2, GPU: Qwen3-0.6B (text) through transformers.js ────────────────────────────────────────────────────
const QWEN = { id: "onnx-community/Qwen3-0.6B-ONNX", revision: "1e0a4a196ecabdf9a879664110574563d3f372d3", dtype: "q4f16", label: "Qwen3-0.6B" };
let tjs = null, qwen = null, qwenLoading = null;
async function loadQwen(progress) {
  if (qwen) return qwen;
  if (qwenLoading) return qwenLoading;
  qwenLoading = (async () => {
    tjs = tjs || (await import("https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0/dist/transformers.min.js"));
    const files = {};
    qwen = await tjs.pipeline("text-generation", QWEN.id, {
      device: gpu.ok ? "webgpu" : "wasm", dtype: gpu.ok ? QWEN.dtype : "int8", revision: QWEN.revision,
      progress_callback: (p) => {
        if (p.status === "progress" && p.total) { files[p.file] = [p.loaded, p.total];
          const [l, t] = Object.values(files).reduce(([a, b], [x, y]) => [a + x, b + y], [0, 0]);
          progress(`${QWEN.label}: downloading ${mb(l)} of ${mb(t)} (once; then from this browser's cache)`); }
        else if (p.status === "ready") progress(`${QWEN.label}: preparing on ${gpu.ok ? "WebGPU" : "the CPU"}…`);
      },
    });
    return qwen;
  })().finally(() => { qwenLoading = null; });
  return qwenLoading;
}
async function askQwen(hist, prompt, push, progress) {
  const gen = await loadQwen(progress);
  const messages = [{ role: "system", content: "You are a helpful, concise assistant running in the visitor's browser on WebGPU. Answer plainly." },
    ...hist.slice(-8), { role: "user", content: `${prompt} /no_think` }];
  const t0 = performance.now();
  let n = 0, first = null, text = "", raw = "", shown = 0;
  progress(`${QWEN.label} on ${gpu.ok ? "WebGPU" : "the CPU"}: reading the prompt…`);
  const streamer = new tjs.TextStreamer(gen.tokenizer, { skip_prompt: true, skip_special_tokens: true,
    callback_function: (t) => {
      if (!t) return;
      if (first === null) first = (performance.now() - t0) / 1000;
      n++; raw += t;
      // Qwen3's reasoning block (empty under /no_think) is not part of the answer: shown neither streaming nor after
      const open = raw.indexOf("<think>"), close = raw.indexOf("</think>");
      if (open !== -1 && close === -1) return;
      const clean = (close !== -1 ? raw.slice(close + 8) : raw).replace(/^\s+/, "");
      if (clean.length > shown) { push(clean.slice(shown)); shown = clean.length; text = clean; }
    } });
  await gen(messages, { max_new_tokens: settings.tokens, do_sample: settings.temp > 0, temperature: settings.temp || undefined, streamer });

  const s = (performance.now() - t0) / 1000;
  return { text, seconds: s, ttft: first, pieces: n, tps: first !== null && s > first ? (n - 1) / (s - first) : null };
}

// ── engine 3, CPU: bankML in this browser (browseML), separate, verified ────────────────────────────────────────
let bankPersona = null, browsed = null;
async function loadBank(progress) {
  if (browseml.isReady()) return browsed;
  if (!bankPersona) bankPersona = await (await fetch("bankml.persona")).json();
  const t0 = performance.now();
  browsed = await browseml.load((got, total, phase) => {
    if (phase === "download") progress(`bankML: downloading ${browseml.MODEL.title}, ${mb(got)} of ${mb(total)} (once)`);
    else if (phase === "cache") progress(`bankML: the model is in this browser (${mb(got)})`);
    else if (phase === "verify") progress(`bankML: checking the model's sha256 against its pin…`);
  }, () => {});
  browsed.loadMs = performance.now() - t0;
  return browsed;
}
async function askBank(hist, prompt, push, progress) {
  const v = await loadBank(progress);
  const where = `WHERE YOU ARE RUNNING (measured by bankML just now): ${v.model} (sha256 ${v.sha256.slice(0, 16)}…), verified by bankML ${v.version}, ` +
    `in the visitor's browser as WebAssembly (browseML) on ${v.threads} CPU thread${v.threads > 1 ? "s" : ""}, beside a WebGPU page; nothing is sent anywhere.`;
  const msgs = [{ role: "system", content: (bankPersona?.system_prompt || "You are bankML.") + "\n\n" + where }, ...hist.slice(-8), { role: "user", content: prompt }];
  const t0 = performance.now();
  let text = "";
  progress(`bankML on ${v.threads} CPU thread${v.threads > 1 ? "s" : ""}: reading the prompt…`);
  const result = await browseml.chat({ messages: msgs, max_tokens: settings.tokens }, (piece) => { text += piece; push(piece); }, () => {});
  const receipt = result.bankml_receipt;
  const ok = !!receipt && (await sha256hex(text)) === receipt.response_sha256;
  return { text, ok, receipt, usage: result.usage || {}, timings: result.timings || {}, seconds: (performance.now() - t0) / 1000, v };
}

// ── one interface over the three, for the field ──────────────────────────────────────────────────────────────
function streamOf(run) {
  const q = []; let wake = null, end = false, err = null;
  const poke = () => { const w = wake; wake = null; w && w(); };
  run((x) => { q.push(x); poke(); }).then(() => { end = true; poke(); }, (e) => { err = e; end = true; poke(); });
  return (async function* () { for (;;) { while (q.length) yield q.shift(); if (end) { if (err) throw err; return; } await new Promise((r) => { wake = r; }); } })();
}
const statusListeners = [];
const which = (mode) => (mode === "persona" ? "bankml" : settings.gpu);
const NAME = {
  get agate() { return `Agate (${gpu.ok ? "WebGPU" : "CPU, no GPU here"})`; },
  get qwen() { return `${QWEN.label} (${gpu.ok ? "WebGPU" : "CPU, no GPU here"})`; },
  bankml: "bankML (CPU)",
};
function showEngines() {
  $("engines").textContent = `generic mode: ${NAME[settings.gpu]} · bankML mode: bankML on the CPU, with a receipt`;
  statusListeners.forEach((f) => f());
}

const engine = {
  async *ask(hist, message, done, _window, progress = () => {}, mode = "chat") {
    const e = which(mode), prompt = message.trim(), date = new Date().toISOString().slice(0, 10);
    let v = null;
    yield* streamOf(async (push) => {
      try {
        if (e === "agate") {
          const r = await askAgate(prompt, progress);
          const url = URL.createObjectURL(new Blob([r.bytes], { type: "image/png" }));
          const sha = await sha256hex(r.bytes);
          push(`🖼 ${prompt}`);
          v = { ok: true, line: `Agate 003 · ${r.width}px · ${secs(r.seconds)}`,
            details: `Agate preview 003 on ${gpu.ok ? `WebGPU (${gpu.name})` : "the CPU"}: ${settings.steps} steps, cfg ${settings.cfg}, seed ${r.seed}, ${r.width}×${r.height} · ${secs(r.seconds)} · marked AI-generated (AGATE003 watermark + PNG fields) · sha256 ${sha.slice(0, 16)}…`,
            media: { kind: "image", url, name: `agate-003-${r.width}px-seed${r.seed}.png`, caption: "AI-generated · Agate preview 003" },
            record: { id: `agate-${Date.now()}`, date, mode: "image", question: prompt, answer: `image/png ${r.width}x${r.height}, ${r.bytes.byteLength} bytes, sha256 ${sha}`,
              carrier: { engine: `onnxruntime-web 1.30.0 ${gpu.ok ? "WebGPU" : "WASM"}`, model: "Logolabs/agate-preview-003", where: `this browser${gpu.ok ? `, ${gpu.name}` : ""}` },
              deliveries: 1, finish: "stop", parameters: { steps: settings.steps, cfg: settings.cfg, seed: r.seed, resolution: settings.res },
              timings: { seconds: r.seconds }, response_sha256: sha, marked_ai_generated: true } };
        } else if (e === "qwen") {
          const r = await askQwen(hist, prompt, push, progress);
          if (!r.text.trim()) push("(no answer: the model wrote nothing)");
          hist.push({ role: "user", content: prompt }, { role: "assistant", content: r.text });
          v = { ok: !!r.text.trim(), line: `${QWEN.label} · ${r.pieces} pieces · ${r.tps ? r.tps.toFixed(1) + " tok/s · " : ""}${secs(r.seconds)}`,
            details: `${QWEN.label} ${gpu.ok ? QWEN.dtype : "int8"} on ${gpu.ok ? `WebGPU (${gpu.name})` : "the CPU"} via transformers.js 4.2.0 · first token ${r.ttft?.toFixed(2) ?? "—"} s · ${r.pieces} pieces · ${secs(r.seconds)} · no receipt (only bankML puts one on an answer)`,
            record: { id: `qwen-${Date.now()}`, date, mode: "text", question: prompt, answer: r.text,
              carrier: { engine: `transformers.js 4.2.0 ${gpu.ok ? "WebGPU" : "WASM"}`, model: `${QWEN.id}@${QWEN.revision.slice(0, 7)} ${gpu.ok ? QWEN.dtype : "int8"}`, where: `this browser${gpu.ok ? `, ${gpu.name}` : ""}` },
              deliveries: 1, finish: "stop", limit: settings.tokens, timings: { ttft_s: r.ttft, seconds: r.seconds, pieces_per_second: r.tps },
              response_sha256: await sha256hex(r.text) } };
        } else {
          const r = await askBank(hist, prompt, push, progress);
          hist.push({ role: "user", content: prompt }, { role: "assistant", content: r.text });
          const u = r.usage, tm = r.timings;
          v = { ok: r.ok, line: `bankML · ${u.completion_tokens ?? "?"} tokens · ${Number(tm.predicted_per_second || 0).toFixed(2)} tok/s · ${r.ok ? "✓ receipt" : "✗ receipt"}`,
            details: `${r.ok ? "✓ receipt: the answer's sha256 matches bankML's" : r.receipt ? "✗ receipt: the answer does not match its sha256" : "✗ no receipt came with this answer"} · ${u.prompt_tokens ?? "?"} + ${u.completion_tokens ?? "?"} tokens counted by bankML ${r.v.version} · ${secs(r.seconds)} · ${r.v.model} (sha256 ${r.v.sha256.slice(0, 12)}…) on ${r.v.threads} CPU thread${r.v.threads > 1 ? "s" : ""}`,
            record: { id: `bankml-${Date.now()}`, date, mode: "bankML", question: prompt, answer: r.text,
              carrier: { engine: `bankML ${r.v.version} browseML (WebAssembly)`, model: r.v.model.replace(/\.gguf$/, ""), sha256: r.v.sha256, system_prompt: "bankml.persona system_prompt", where: "this browser (browseML)" },
              deliveries: 1, finish: "stop", limit: settings.tokens, usage: u,
              timings: { ttft_s: tm.prompt_ms != null ? tm.prompt_ms / 1000 : undefined, seconds: r.seconds, predicted_per_second: tm.predicted_per_second },
              response_sha256: await sha256hex(r.text), receipt: { ok: r.ok } } };
        }
      } catch (err) {
        const m = `${NAME[e]} did not answer: ${String(err?.message || err).slice(0, 300)}`;
        push(m);
        v = { ok: false, line: m };
      }
    });
    done(v || { ok: false, line: "no answer" });
  },
  status() { return `${NAME[settings.gpu]} · bankML ${browseml.isReady() ? "✓" : "—"}`; },
  onStatus(f) { statusListeners.push(f); },
  async ping() {
    return { ok: gpu.ok, line: `${gpu.ok ? `WebGPU: ${gpu.name}${gpu.f16 ? ", shader-f16" : ""}, max buffer ${mb(gpu.maxBuffer)}` : `no WebGPU: ${gpu.why}`} · Agate ${agate ? "loaded" : "not loaded"} · ${QWEN.label} ${qwen ? "loaded" : "not loaded"} · bankML ${browseml.isReady() ? `loaded, ${browsed?.threads} thread(s)` : "not loaded"}` };
  },
};

// ── the button set; every one is also a T-mode command ───────────────────────────────────────────────────────
const num = (x, lo, hi) => { const n = Number(x); return Number.isFinite(n) && n >= lo && n <= hi ? n : undefined; };
const setting = (cmd, key, lo, hi, int, what) => ({
  id: `wg-${cmd}`, label: cmd, command: cmd, icon: "⚙", defaultHidden: true, description: `${cmd} <${lo}-${hi}> - ${what}`,
  run(ctx) {
    if (key === "seed" && ctx.args[0] === "random") { settings.seed = null; save(); return ctx.print("seed: random"); }
    const n = num(ctx.args[0], lo, hi);
    if (n === undefined) return ctx.print(`${cmd}: now ${settings[key] ?? "random"}; give ${lo} to ${hi}${key === "seed" ? " or random" : ""}`, "err");
    settings[key] = int ? Math.round(n) : n; save(); ctx.print(`${cmd}: ${settings[key]}`);
    if (key === "res" && agate) { agate = null; ctx.print("Agate reloads at the new resolution on the next image"); }
  },
});
engine.actions = [
  { id: "wg-gpu", label: "GPU engine: Qwen3 text or Agate images", command: "gpu", aliases: ["engine"], icon: "⧉",
    description: "gpu [qwen|agate] - what the generic mode runs on WebGPU (no argument: switch)",
    run(ctx) {
      const a = ctx.args[0];
      settings.gpu = a === "agate" || a === "image" ? "agate" : a === "qwen" || a === "text" ? "qwen" : settings.gpu === "qwen" ? "agate" : "qwen";
      save(); ctx.print(`generic mode runs ${NAME[settings.gpu]}`);
    } },
  { id: "wg-load", label: "Load the models", command: "load", icon: "⬇",
    description: "load [agate|qwen|bankml|all] - download (once) and prepare a model now, instead of on the first question",
    async run(ctx) {
      const w = ctx.args[0] || (ctx.mode === "persona" ? "bankml" : settings.gpu);
      const list = w === "all" ? ["agate", "qwen", "bankml"] : [w];
      for (const k of list) {
        const t0 = performance.now(), say = (t) => {};
        try {
          if (k === "agate") await loadAgate(say); else if (k === "qwen") await loadQwen(say); else await loadBank(say);
          ctx.print(`✓ ${NAME[k]} ready in ${secs((performance.now() - t0) / 1000)}`);
        } catch (e) { ctx.print(`${NAME[k]} could not load: ${String(e?.message || e).slice(0, 200)}`, "err"); }
      }
      showEngines();
    } },
  { id: "wg-example", label: "An example", command: "example", aliases: ["ex"], icon: "✦",
    description: "example [n] - an example for the engine now (no number: list them)",
    run(ctx) {
      const ex = ctx.mode === "persona" ? ["Who are you, and how do I know this answer is yours?", "What is a receipt on an answer?"]
        : settings.gpu === "agate" ? ["A lighthouse on a cliff at dusk, oil painting", "A red panda reading a book in a library, watercolor", "A neon sign that says \"WEBGPU\" on a brick wall"]
        : ["Explain WebGPU in three sentences.", "Write a haiku about a GPU.", "What is the difference between a CPU and a GPU?"];
      const n = Number(ctx.args[0]);
      if (!n || !ex[n - 1]) { ex.forEach((e, i) => ctx.print(`${i + 1}. ${e}`)); return; }
      ctx.send(ex[n - 1]);
    } },
  setting("tokens", "tokens", 16, 2048, true, "the most a text answer may use"),
  setting("temp", "temp", 0, 1.5, false, "Qwen3's temperature (0 answers the same each time)"),
  setting("steps", "steps", 4, 100, true, "Agate's sampling steps (fewer is faster)"),
  setting("cfg", "cfg", 1, 10, false, "Agate's guidance scale"),
  setting("res", "res", 256, 512, true, "Agate's resolution: 512 or 256"),
  setting("seed", "seed", 0, 2 ** 31 - 1, true, "Agate's seed (random: a new one each image)"),
  { id: "wg-help", label: "WebGPU commands", command: "webgpu", aliases: ["wg"], icon: "⧉", defaultHidden: true, description: "webgpu - these commands",
    run(ctx) { for (const a of engine.actions) ctx.print(`  ${a.description}`); ctx.print("  bankML - the bankML mode · > the generic mode · ping · last · details · help"); } },
];

// ── boot ─────────────────────────────────────────────────────────────────────────────────────────────────────
showEngines();
checkGpu().then(showEngines);
(() => {
  const ui = window.BankmlSpaceUIF;
  if (!ui) return;
  if (ui.isOutputWindow()) document.body.classList.add("uif-popout");
  ui.mount($("uif-root"), engine, {
    placeholder: "Ask, or describe an image (⧉ switches Qwen3 / Agate) · bankML for the verified CPU engine · T terminal",
    terminalPlaceholder: "T mode: webgpu · gpu · load all · steps · seed · tokens · ping · last · help",
    anchor: $("top"), storageKey: "ultimate-webgpu",
    persona: { label: "bankML", title: "bankML mode: bankML answers on your CPU, verified, with a receipt",
      placeholder: "bankML: verified answers on your CPU, with a receipt (Enter) · > the GPU · T terminal" },
    pill: { idle: "no answer yet", ok: "✓ done", bad: "✗ not done" },
    modeBar: {
      chat: ["wg-gpu", "wg-load", "wg-example", "spawn", "outputs"],
      persona: ["wg-load", "wg-example", "spawn", "outputs"],
      terminal: ["wg-help", "wg-gpu", "wg-load", "last", "details", "ping", "outputs", "clear"],
    },
  });
})();