File size: 9,916 Bytes
520f98d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
<!doctype html>
<meta charset="utf-8" />
<title>MangaOCR WebGPU comparison harness</title>
<style>
	body { background: #111; color: #eee; font: 14px/1.45 ui-monospace, monospace; margin: 24px; }
	pre { white-space: pre-wrap; }
</style>
<h1>MangaOCR WebGPU comparison harness</h1>
<pre id="log">Starting…</pre>
<script type="module">
	import * as ort from "./.cache/ort/ort.webgpu.min.mjs";

	const parameters = new URLSearchParams(location.search);
	const output = document.querySelector("#log");
	const lines = [];
	const log = (message) => {
		lines.push(`${new Date().toISOString()} ${message}`);
		output.textContent = lines.join("\n");
		console.log(message);
	};
	const timed = async (operation) => {
		const started = performance.now();
		const value = await operation();
		return [value, performance.now() - started];
	};
	window.addEventListener("error", (event) => log(`FAIL ${event.error?.stack ?? event.message}`));
	window.addEventListener("unhandledrejection", (event) => log(`FAIL ${event.reason?.stack ?? event.reason}`));
	if (!navigator.gpu) throw new Error("WebGPU is unavailable");
	log(`navigator.gpu=true secure=${window.isSecureContext}`);

	const session = async (path) => {
		const bytes = new Uint8Array(await (await fetch(path)).arrayBuffer());
		return ort.InferenceSession.create(bytes, {
			executionProviders: ["webgpu"],
			logSeverityLevel: 3,
		});
	};
	let value;
	let elapsed;
	[value, elapsed] = await timed(() => session("./.cache/mangaocr/encoder_model.onnx"));
	const encoder = value;
	log(`PASS encoder session creation (${elapsed.toFixed(1)} ms)`);
	[value, elapsed] = await timed(() => session("./.cache/mangaocr/decoder_model.onnx"));
	const decoder = value;
	log(`PASS decoder session creation (${elapsed.toFixed(1)} ms)`);
	const vocab = (await (await fetch("./.cache/mangaocr/vocab.txt")).text()).split(/\r?\n/g);
	log(`models ready vocab=${vocab.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];
		return raw.split(",").map(Number);
	};
	const cropImageData = (image) => {
		const [x, y, width, height] = parseCrop(image);
		const canvas = document.createElement("canvas");
		canvas.width = width;
		canvas.height = height;
		const context = canvas.getContext("2d", { willReadFrequently: true });
		context.drawImage(image, x, y, width, height, 0, 0, width, height);
		return context.getImageData(0, 0, width, height);
	};
	const clampByte = (number) => Math.max(0, Math.min(255, number));
	const grayscale = (image) => {
		const result = new Uint8ClampedArray(image.width * image.height);
		for (let source = 0, pixel = 0; source < image.data.length; source += 4, pixel += 1) {
			const r = image.data[source] ?? 0;
			const g = image.data[source + 1] ?? 0;
			const b = image.data[source + 2] ?? 0;
			result[pixel] = ((b * 19595 + g * 38470 + r * 7471 + 0x8000) >> 16) & 0xff;
		}
		return result;
	};
	const coefficients = (inputSize, outputSize) => {
		const scale = inputSize / outputSize;
		const support = scale >= 1 ? scale : 1;
		const maximum = Math.ceil(support) * 2 + 1;
		const result = [];
		for (let outputIndex = 0; outputIndex < outputSize; outputIndex += 1) {
			const center = (outputIndex + 0.5) * scale;
			const inverse = scale >= 1 ? 1 / scale : 1;
			const start = Math.max(0, Math.min(inputSize, Math.trunc(center - support + 0.5)));
			const end = Math.min(Math.trunc(center + support + 0.5), inputSize);
			const size = Math.max(0, Math.min(maximum, end - start));
			const weights = [];
			let total = 0;
			for (let index = 0; index < size; index += 1) {
				const weight = Math.max(0, 1 - Math.abs((index + start - center + 0.5) * inverse));
				weights.push(weight);
				total += weight;
			}
			result.push({
				start,
				size,
				weights: weights.map((weight) => Math.trunc(0.5 + weight / total * 2 ** 22)),
			});
		}
		return result;
	};
	const preprocess = (image) => {
		const source = grayscale(image);
		const horizontal = coefficients(image.width, 224);
		const vertical = coefficients(image.height, 224);
		const temporary = new Uint8ClampedArray(image.height * 224);
		for (let y = 0; y < image.height; y += 1) {
			for (let x = 0; x < 224; x += 1) {
				const item = horizontal[x];
				let sum = 1 << 21;
				for (let index = 0; index < item.size; index += 1) {
					sum += (source[y * image.width + item.start + index] ?? 0) * (item.weights[index] ?? 0);
				}
				temporary[y * 224 + x] = clampByte(Math.trunc(sum / 2 ** 22));
			}
		}
		const resized = new Uint8ClampedArray(224 * 224);
		for (let y = 0; y < 224; y += 1) {
			const item = vertical[y];
			for (let x = 0; x < 224; x += 1) {
				let sum = 1 << 21;
				for (let index = 0; index < item.size; index += 1) {
					sum += (temporary[(item.start + index) * 224 + x] ?? 0) * (item.weights[index] ?? 0);
				}
				resized[y * 224 + x] = clampByte(Math.trunc(sum / 2 ** 22));
			}
		}
		const plane = 224 * 224;
		const result = new Float32Array(3 * plane);
		for (let pixel = 0; pixel < plane; pixel += 1) {
			const normalized = Math.fround(Math.fround(Math.fround(resized[pixel]) / Math.fround(255)) * Math.fround(2) - Math.fround(1));
			result[pixel] = normalized;
			result[plane + pixel] = normalized;
			result[2 * plane + pixel] = normalized;
		}
		return result;
	};
	const logSoftmax = (values) => {
		let maximum = Number.NEGATIVE_INFINITY;
		for (const item of values) maximum = Math.max(maximum, item);
		let sum = 0;
		for (const item of values) sum += Math.exp(item - maximum);
		const logZ = Math.log(sum) + maximum;
		return Float32Array.from(values, (item) => item - logZ);
	};
	const noRepeat = (ids, scores) => {
		if (ids.length < 2) return;
		const prefix = ids.slice(-2);
		for (let index = 0; index <= ids.length - 3; index += 1) {
			if (prefix.every((id, offset) => id === ids[index + offset])) scores[ids[index + 2]] = Number.NEGATIVE_INFINITY;
		}
	};
	const top = (items, item, count) => {
		if (items.length < count || item.score > (items.at(-1)?.score ?? Number.NEGATIVE_INFINITY)) {
			items.push(item);
			items.sort((a, b) => b.score - a.score);
			if (items.length > count) items.pop();
		}
	};
	const normalizedScore = (beam) => beam.score / Math.pow(Math.max(1, beam.ended ? beam.ids.length - 1 : beam.ids.length), 2);
	const decodeText = (ids) => ids.filter((id) => id >= 15).map((id) => vocab[id] ?? "").map((token) => token.replace(/^##/, "")).join("").replace(/\s+/g, "").replaceAll("…", "...").replace(/[・.]{2,}/g, (match) => ".".repeat(match.length)).replace(/[\u0021-\u007e]/g, (character) => String.fromCharCode(character.charCodeAt(0) + 0xfee0));

	const recognize = async (image) => {
		const pixels = preprocess(cropImageData(image));
		let result;
		let encoderMs;
		[result, encoderMs] = await timed(() => encoder.run({ pixel_values: new ort.Tensor("float32", pixels, [1, 3, 224, 224]) }));
		const hiddenTensor = result.last_hidden_state ?? Object.values(result)[0];
		log(`encoder output=${hiddenTensor.location} ${JSON.stringify(hiddenTensor.dims)}`);
		const hidden = hiddenTensor.data;
		const [, seqLen, hiddenDim] = hiddenTensor.dims;
		let active = [{ ids: [2], score: 0, ended: false }];
		const completed = [];
		const decoderStarted = performance.now();
		for (let step = 1; step < 96; step += 1) {
			if (!active.length || completed.length >= 4) break;
			const tokenLength = active[0].ids.length;
			const inputIds = new BigInt64Array(active.length * tokenLength);
			const repeatedHidden = new Float32Array(active.length * seqLen * hiddenDim);
			active.forEach((beam, beamIndex) => {
				beam.ids.forEach((id, idIndex) => inputIds[beamIndex * tokenLength + idIndex] = BigInt(id));
				repeatedHidden.set(hidden, beamIndex * seqLen * hiddenDim);
			});
			const decoderResult = await decoder.run({
				input_ids: new ort.Tensor("int64", inputIds, [active.length, tokenLength]),
				encoder_hidden_states: new ort.Tensor("float32", repeatedHidden, [active.length, seqLen, hiddenDim]),
			});
			const tensor = decoderResult.logits ?? Object.values(decoderResult)[0];
			const candidates = [];
			active.forEach((beam, beamIndex) => {
				const rowOffset = beamIndex * tokenLength * 6144 + (tokenLength - 1) * 6144;
				const scores = logSoftmax(tensor.data.subarray(rowOffset, rowOffset + 6144));
				noRepeat(beam.ids, scores);
				for (let token = 0; token < scores.length; token += 1) top(candidates, { beam, token, score: beam.score + scores[token] }, 8);
			});
			const next = [];
			for (let rank = 0; rank < candidates.length && next.length < 4; rank += 1) {
				const candidate = candidates[rank];
				if (candidate.token === 3) {
					if (rank < 4) completed.push({ ids: [...candidate.beam.ids, 3], score: candidate.score, ended: true });
				} else next.push({ ids: [...candidate.beam.ids, candidate.token], score: candidate.score, ended: false });
			}
			active = next;
		}
		const decoderMs = performance.now() - decoderStarted;
		const candidates = (completed.length >= 4 ? completed : [...completed, ...active]).sort((a, b) => normalizedScore(b) - normalizedScore(a));
		const text = decodeText(candidates[0]?.ids.filter((id) => id !== 2 && id !== 3) ?? []);
		hiddenTensor.dispose();
		return { encoderMs, decoderMs, totalMs: encoderMs + decoderMs, text };
	};

	const image = await loadImage(parameters.get("image") ?? "/.cache/demo/daf0244b038a-20260706.png");
	const runs = Number(parameters.get("runs") ?? 4);
	for (let run = 1; run <= runs; run += 1) {
		const result = await recognize(image);
		log(`RUN ${run}/${runs} encoder=${result.encoderMs.toFixed(1)} decoder=${result.decoderMs.toFixed(1)} total=${result.totalMs.toFixed(1)} ms`);
		log(`TEXT ${result.text}`);
	}
</script>