Spaces:

ML-Intern-lab
/

agate-preview-002-4step-webgpu

Running

agate-preview-002-4step-webgpu / src /js /app.js

Move to ML-Intern-lab: new repo names

eb5f3ed verified 24 days ago

6.79 kB

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 530 MB)"; };
      window.__agate = { state: "no-webgpu", why: s.why };
      return;
    }
    ui.statusR.textContent = s.adapter || "";
  }
  if (params.has("autoload") || (await hasCachedModel(VERSION))) load();
}

init();

	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 530 MB)"; };
	window.__agate = { state: "no-webgpu", why: s.why };
	return;
	}
	ui.statusR.textContent = s.adapter \|\| "";
	}
	if (params.has("autoload") \|\| (await hasCachedModel(VERSION))) load();
	}

	init();