vision-worker.mjs

import * as ort from "./.work/ort/ort-webgpu.mjs";

let session = null;
let device = null;

const reply = (requestId, payload, transfer = []) => postMessage({ requestId, payload }, transfer);
const fail = (requestId, error) => postMessage({ requestId, error: error?.stack ?? String(error) });

self.onmessage = async ({ data }) => {
	const { requestId, type, payload } = data;
	try {
		if (type === "init") {
			ort.env.wasm.numThreads = 1;
			ort.env.wasm.wasmPaths = payload.wasmPaths;
			const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
			if (!adapter) throw new Error("No WebGPU adapter is available in the vision worker");
			device = await adapter.requestDevice();
			ort.env.webgpu.device = device;
			const options = {
				executionProviders: ["webgpu"],
				preferredOutputLocation: { vision_embeds: "cpu" },
				logSeverityLevel: 3,
				...(payload.lean ? { enableCpuMemArena: false, enableMemPattern: false } : {}),
			};
			session = await ort.InferenceSession.create(payload.modelPath, options);
			if (JSON.stringify(session.inputNames) !== JSON.stringify(["pixel_values"])) {
				throw new Error(`Unexpected vision inputs ${session.inputNames}`);
			}
			reply(requestId, { ready: true });
			return;
		}
		if (type === "encode") {
			if (!session) throw new Error("Vision worker is not initialized");
			const pixels = new ort.Tensor("float32", payload.pixels, [1, 3, 224, 224]);
			const started = performance.now();
			const result = await session.run({ pixel_values: pixels });
			const visionMs = performance.now() - started;
			pixels.dispose();
			const values = new Float32Array(result.vision_embeds.data);
			const dims = [...result.vision_embeds.dims];
			result.vision_embeds.dispose();
			reply(requestId, { values, dims, visionMs }, [values.buffer]);
			return;
		}
		if (type === "close") {
			if (session) await session.release();
			session = null;
			if (device) device.destroy();
			device = null;
			reply(requestId, { closed: true });
			return;
		}
		throw new Error(`Unknown worker operation ${type}`);
	} catch (error) {
		fail(requestId, error);
	}
};