| <!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> |
|
|
