webgpu-e2e.html

<!doctype html>
<meta charset="utf-8" />
<title>Baberu end-to-end WebGPU OCR</title>
<style>
	body { background: #111; color: #eee; font: 14px/1.45 ui-monospace, monospace; margin: 24px; }
	.controls { display: flex; gap: 12px; align-items: center; margin-bottom: 16px; }
	button, input { font: inherit; }
	canvas { border: 1px solid #555; height: 224px; width: 224px; }
	pre { white-space: pre-wrap; }
</style>
<h1>Baberu end-to-end WebGPU OCR</h1>
<div class="controls">
	<input id="file" type="file" accept="image/*" />
	<button id="run" type="button">Run selected crop</button>
</div>
<canvas id="preview" width="224" height="224"></canvas>
<pre id="log">Starting…</pre>
<script type="module">
	import * as ort from "./.cache/ort/ort.webgpu.min.mjs";

	const VOCAB_SIZE = 14630;
	const BOS = 1;
	const EOS = 2;
	const MEAN = [0.485, 0.456, 0.406];
	const STD = [0.229, 0.224, 0.225];
	const cacheNames = [
		...Array.from({ length: 6 }, (_, index) => `present_k${index}`),
		...Array.from({ length: 6 }, (_, index) => `present_v${index}`),
	];
	const output = document.querySelector("#log");
	const preview = document.querySelector("#preview");
	const context = preview.getContext("2d", { willReadFrequently: true });
	const lines = [];
	const log = (message) => {
		lines.push(`${new Date().toISOString()} ${message}`);
		output.textContent = lines.join("\n");
		console.log(message);
	};
	const timed = async (label, operation) => {
		const started = performance.now();
		const result = await operation();
		log(`PASS ${label} (${(performance.now() - started).toFixed(1)} ms)`);
		return result;
	};
	window.addEventListener("error", (event) => log(`FAIL ${event.error?.stack ?? event.message}`));
	window.addEventListener("unhandledrejection", (event) => log(`FAIL ${event.reason?.stack ?? event.reason}`));

	ort.env.wasm.numThreads = 1;
	ort.env.wasm.wasmPaths = "/.cache/ort/";
	if (!navigator.gpu) throw new Error("WebGPU is unavailable");
	const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
	if (!adapter) throw new Error("No WebGPU adapter is available");
	const gpuDevice = await adapter.requestDevice();
	ort.env.webgpu.device = gpuDevice;
	log(`navigator.gpu=true secure=${window.isSecureContext}`);

	const createSession = async (path, preferredOutputLocation) => {
		const response = await fetch(path);
		if (!response.ok) throw new Error(`${path}: HTTP ${response.status}`);
		const bytes = new Uint8Array(await response.arrayBuffer());
		return ort.InferenceSession.create(bytes, {
			executionProviders: ["webgpu"],
			preferredOutputLocation,
			logSeverityLevel: 2,
		});
	};

	const parameters = new URLSearchParams(location.search);
	const bundle = parameters.get("bundle") ?? "lossless";
	const compactVision = bundle.startsWith("compact-");
	const compact = bundle === "compact-qdq";
	const compactFp16 = bundle === "compact-fp16";
	const gatherTokens = bundle === "compact-unified-gather" || bundle === "balanced-unified-gather";
	const unifiedQdq = bundle === "compact-unified-qdq" || bundle === "balanced-unified-qdq" || gatherTokens;
	const qdq = compact || bundle === "balanced-qdq";
	const publishedLayout = parameters.get("layout") === "hf";
	const publishedRoot = compactVision ? "./variants/webgpu-121" : "./variants/webgpu-242";
	const visionPath = publishedLayout
		? `${publishedRoot}/${compactVision ? "vision_int4.onnx" : "vision_fp16.onnx"}`
		: compactVision
			? "./model/onnx/vision_int4.onnx"
			: parameters.get("vision") === "fp32"
				? "./output/vision_fp32.onnx"
				: "./model/onnx/vision_fp16.onnx";
	const prefillPath = publishedLayout
		? `${publishedRoot}/decoder_prefill_qdq_int8.onnx`
		: compactFp16
			? "./output/decoder_prefill_fp16.onnx"
			: qdq
				? "./output/decoder_prefill_qdq_int8.onnx"
				: "./output/decoder_prefill_fp32.onnx";
	const stepPath = publishedLayout
		? `${publishedRoot}/decoder_step_qdq_int8.onnx`
		: compactFp16
			? "./output/decoder_step_fp16.onnx"
			: qdq
				? "./output/decoder_step_qdq_int8.onnx"
				: "./output/decoder_step_fp32.onnx";
	const maxTokens = Number(parameters.get("tokens") ?? 128);
	const vision = await timed("vision session creation", () =>
		createSession(visionPath, { vision_embeds: "gpu-buffer" })
	);
	const decoderOutputs = Object.fromEntries([
		["logits", "cpu"],
		...cacheNames.map((name) => [name, "gpu-buffer"]),
	]);
	const unifiedFile = gatherTokens
		? "decoder_unified_gather_qdq_int8.onnx"
		: "decoder_unified_qdq_int8.onnx";
	const unifiedPath = publishedLayout
		? `${publishedRoot}/${unifiedFile}`
		: `./output/${unifiedFile}`;
	const activePrefillPath = unifiedQdq ? unifiedPath : prefillPath;
	const activeStepPath = unifiedQdq ? unifiedPath : stepPath;
	const prefill = await timed("prefill session creation", () =>
		createSession(activePrefillPath, decoderOutputs)
	);
	const step = unifiedQdq ? prefill : await timed("step session creation", () =>
		createSession(stepPath, decoderOutputs)
	);
	let sessionsReleased = false;
	const releaseSessions = async () => {
		if (sessionsReleased) return;
		sessionsReleased = true;
		const sessions = step === prefill ? [vision, prefill] : [vision, prefill, step];
		const results = await Promise.allSettled(sessions.map((session) => session.release()));
		gpuDevice.destroy();
		const failed = results.filter((result) => result.status === "rejected");
		if (failed.length) {
			throw new Error(`Failed to release ${failed.length}/${sessions.length} model sessions`);
		}
		document.querySelector("#run").disabled = true;
		log(`PASS released ${sessions.length} model sessions`);
	};
	window.releaseSessions = releaseSessions;
	window.addEventListener("pagehide", () => void releaseSessions(), { once: true });
	const vocabPath = publishedLayout ? "./tokenizer/vocab.json" : "./model/tokenizer/vocab.json";
	const charset = await (await fetch(vocabPath)).json();
	const idToCharacter = ["", "", "", "", ...charset];
	const contentIds = new Set();
	for (let index = 0; index < charset.length; index += 1) {
		const character = charset[index];
		if (character.length === 1 && !"ーｰ〜~".includes(character) && /[\p{Letter}\p{Number}]/u.test(character)) {
			contentIds.add(index + 4);
		}
	}
	log(`models ready bundle=${bundle} vision=${visionPath} prefill=${activePrefillPath} step=${activeStepPath} sharedDecoder=${unifiedQdq} gatherTokens=${gatherTokens} vocab=${idToCharacter.length}`);

	const loadImage = (source) => new Promise((resolve, reject) => {
		const image = new Image();
		image.onload = () => resolve(image);
		image.onerror = reject;
		image.src = source;
	});
	const parseCrop = (image) => {
		const raw = parameters.get("crop");
		if (!raw) return [0, 0, image.naturalWidth, image.naturalHeight];
		const values = raw.split(",").map(Number);
		if (values.length !== 4 || values.some((value) => !Number.isFinite(value))) {
			throw new Error("crop must be x,y,width,height");
		}
		return values;
	};
	const preprocess = (image) => {
		const crop = parseCrop(image);
		context.imageSmoothingEnabled = true;
		context.imageSmoothingQuality = "high";
		context.clearRect(0, 0, 224, 224);
		context.drawImage(image, ...crop, 0, 0, 224, 224);
		const rgba = context.getImageData(0, 0, 224, 224).data;
		const values = new Float32Array(3 * 224 * 224);
		for (let pixel = 0; pixel < 224 * 224; pixel += 1) {
			for (let channel = 0; channel < 3; channel += 1) {
				values[channel * 224 * 224 + pixel] =
					(rgba[pixel * 4 + channel] / 255 - MEAN[channel]) / STD[channel];
			}
		}
		log(`crop=${crop.join(",")} source=${image.naturalWidth}x${image.naturalHeight}`);
		return new ort.Tensor("float32", values, [1, 3, 224, 224]);
	};
	const oneHot = (token) => {
		const values = new Float32Array(VOCAB_SIZE);
		values[token] = 1;
		return new ort.Tensor("float32", values, [1, 1, VOCAB_SIZE]);
	};
	const tokenId = (token) => new ort.Tensor("int32", Int32Array.of(token), [1, 1]);
	const emptyVision = () => new ort.Tensor("float32", new Float32Array(0), [1, 0, 512]);
	const emptyCache = () => new ort.Tensor("float32", new Float32Array(0), [1, 2, 0, 64]);
	const prefillPositions = () => new ort.Tensor(
		"int32",
		Int32Array.from({ length: 257 }, (_, index) => index),
		[1, 257],
	);
	const shouldBlockToken = (tokens) => {
		if (!tokens.length || !contentIds.has(tokens.at(-1))) return false;
		const last = tokens.at(-1);
		let run = 0;
		for (let index = tokens.length - 1; index >= 0 && tokens[index] === last; index -= 1) run += 1;
		return run >= 12;
	};
	const chooseToken = (source, sequence, tokens) => {
		const seen = new Set(sequence);
		const blocked = shouldBlockToken(tokens) ? tokens.at(-1) : -1;
		const adjusted = (index) => {
			if (index === blocked) return Number.NEGATIVE_INFINITY;
			const value = source[index];
			if (!seen.has(index)) return value;
			return value < 0 ? value * 1.2 : value / 1.2;
		};
		let result = 0;
		let best = adjusted(0);
		for (let index = 1; index < source.length; index += 1) {
			const value = adjusted(index);
			if (value > best) {
				best = value;
				result = index;
			}
		}
		return result;
	};
	const disposeResult = (result) => {
		for (const name of cacheNames) result[name].dispose();
		result.logits.dispose();
	};

	const recognize = async (image) => {
		const pixels = preprocess(image);
		const totalStarted = performance.now();
		const visionStarted = performance.now();
		const visionResult = await timed("vision execution", () => vision.run({ pixel_values: pixels }));
		const visionMs = performance.now() - visionStarted;
		log(`vision output=${visionResult.vision_embeds.location} ${JSON.stringify(visionResult.vision_embeds.dims)}`);
		const prefillStarted = performance.now();
		const prefillFeeds = unifiedQdq
			? {
				vision_embeds: visionResult.vision_embeds,
				[gatherTokens ? "token_ids" : "token_one_hot"]: gatherTokens ? tokenId(BOS) : oneHot(BOS),
				position_ids: prefillPositions(),
				...Object.fromEntries(cacheNames.map((name) => [name.replace("present_", "past_"), emptyCache()])),
			}
			: { vision_embeds: visionResult.vision_embeds };
		let cacheResult = await timed("decoder prefill", () => prefill.run(prefillFeeds));
		if (unifiedQdq) {
			prefillFeeds[gatherTokens ? "token_ids" : "token_one_hot"].dispose();
			prefillFeeds.position_ids.dispose();
			for (const name of cacheNames) prefillFeeds[name.replace("present_", "past_")].dispose();
		}
		const prefillMs = performance.now() - prefillStarted;
		log(`KV output=${cacheResult.present_k0.location} ${JSON.stringify(cacheResult.present_k0.dims)}`);
		visionResult.vision_embeds.dispose();
		pixels.dispose();

		const sequence = [BOS];
		const tokens = [];
		const decodeStarted = performance.now();
		for (let iteration = 0; iteration < maxTokens; iteration += 1) {
			const next = chooseToken(cacheResult.logits.data, sequence, tokens);
			if (next === EOS) break;
			tokens.push(next);
			sequence.push(next);
			if (tokens.length >= maxTokens) break;
			const feeds = {
				[gatherTokens ? "token_ids" : "token_one_hot"]: gatherTokens ? tokenId(next) : oneHot(next),
				position_ids: new ort.Tensor("int32", new Int32Array([257 + iteration]), [1, 1]),
			};
			if (unifiedQdq) feeds.vision_embeds = emptyVision();
			for (let layer = 0; layer < 6; layer += 1) {
				feeds[`past_k${layer}`] = cacheResult[`present_k${layer}`];
				feeds[`past_v${layer}`] = cacheResult[`present_v${layer}`];
			}
			const nextResult = await step.run(feeds);
			disposeResult(cacheResult);
			feeds[gatherTokens ? "token_ids" : "token_one_hot"].dispose();
			feeds.position_ids.dispose();
			if (unifiedQdq) feeds.vision_embeds.dispose();
			cacheResult = nextResult;
		}
		const decodeMs = performance.now() - decodeStarted;
		const text = tokens.map((token) => idToCharacter[token] ?? "").join("");
		log(`PASS decode ${tokens.length} tokens (${decodeMs.toFixed(1)} ms)`);
		log(`TEXT ${text}`);
		disposeResult(cacheResult);
		return {
			text,
			tokens: tokens.length,
			visionMs,
			prefillMs,
			decodeMs,
			totalMs: performance.now() - totalStarted,
		};
	};
	const normalizeText = (text) => text.normalize("NFKC").replace(/\s+/g, "");
	const editDistance = (left, right) => {
		let previous = Array.from({ length: right.length + 1 }, (_, index) => index);
		for (let leftIndex = 1; leftIndex <= left.length; leftIndex += 1) {
			const current = [leftIndex];
			for (let rightIndex = 1; rightIndex <= right.length; rightIndex += 1) {
				current.push(Math.min(
					current[rightIndex - 1] + 1,
					previous[rightIndex] + 1,
					previous[rightIndex - 1] + (left[leftIndex - 1] === right[rightIndex - 1] ? 0 : 1),
				));
			}
			previous = current;
		}
		return previous.at(-1);
	};
	const hayai13 = [
		["01", "知らない世界で見つけたイメージを"],
		["02", "カナデトモスソラ（Ｋａｎａｄｅｔｏｍｏｓｕｓｏｒａ）"],
		["03", "建設会社社員行方"],
		["04", "だとしてもこのレベルがウロつくなんて．．．おそらく２級の呪い"],
		["05", "パチパチパチパチ"],
		["06", "バビュン"],
		["07", "僕の過去とか未来とか"],
		["08", "くらべられっ子"],
		["09", "そうだクラス分けがあるんだった！！"],
		["10", "脇役よ、主役を超えよ！"],
		["11", "Ｅｈ～Ｉｄｏｎ＇ｔｒｅａｌｌｙｗａｎｔｔｏ～"],
		["12", "「Ｓｏｒｒｙｆｏｒｔｈｅｗａｉｔ～！Ｄｉｄｙｏｕｗａｉｔｌｏｎｇ？」"],
		["13", "ＹａｍａｔｅＡｒｅａＮｅｗｒｅｓｉｄｅｎｔｉａｌｄｉｓｔｒｉｃｔｆｏｒｆｏｒｅｉｇｎｅｒｓ"],
	];
	const runHayai13 = async () => {
		const results = [];
		for (const [id, expected] of hayai13) {
			log(`CASE ${id}/13`);
			const result = await recognize(await loadImage(`./.cache/hayai13/${id}.png`));
			const normalizedExpected = normalizeText(expected);
			const normalizedActual = normalizeText(result.text);
			results.push({
				id,
				expected,
				actual: result.text,
				distance: editDistance(normalizedExpected, normalizedActual),
				characters: normalizedExpected.length,
				exact: normalizedExpected === normalizedActual,
				...result,
			});
		}
		const distance = results.reduce((sum, result) => sum + result.distance, 0);
		const characters = results.reduce((sum, result) => sum + result.characters, 0);
		const exact = results.filter((result) => result.exact).length;
		const sortedTimes = results.map((result) => result.totalMs).sort((a, b) => a - b);
		const summary = {
			bundle,
			cases: results.length,
			nCER: distance / characters,
			distance,
			characters,
			exact,
			medianMs: sortedTimes[Math.floor(sortedTimes.length / 2)],
			peakCaseMs: sortedTimes.at(-1),
			results,
		};
		log(`SUITE_RESULT ${JSON.stringify(summary)}`);
		window.__suiteResult = summary;
		await fetch("/benchmark-result", {
			method: "POST",
			headers: { "Content-Type": "application/json" },
			body: JSON.stringify(summary),
		});
		log(`PASS persisted benchmark result for ${bundle}`);
	};
	window.runHayai13 = runHayai13;

	const fileInput = document.querySelector("#file");
	document.querySelector("#run").addEventListener("click", async () => {
		const [file] = fileInput.files;
		if (!file) throw new Error("Choose a speech-bubble crop first");
		try {
			await recognize(await loadImage(URL.createObjectURL(file)));
		} finally {
			await releaseSessions();
		}
	});
	const suite = parameters.get("suite");
	const imagePath = parameters.get("image");
	if (suite === "hayai13" || imagePath) {
		try {
			if (suite === "hayai13") {
				await runHayai13();
			} else {
				const image = await loadImage(imagePath);
				const runs = Number(parameters.get("runs") ?? 1);
				for (let run = 1; run <= runs; run += 1) {
					log(`RUN ${run}/${runs}`);
					await recognize(image);
				}
			}
		} finally {
			await releaseSessions();
		}
	}
</script>