Add host ONNX runtime option

src/index.js

@@ -40,6 +40,7 @@ async function createStructuredDocumentText(buf, options = {}) {
 		contentType,
 		password,
 		dataProvider,
+		nativeONNXRun,
 		sourceHash,
 	} = options;
 	assertSourceHash(sourceHash);
@@ -50,6 +51,7 @@ async function createStructuredDocumentText(buf, options = {}) {
 			: await getSnapshotStructure(buf, contentType, { sourceHash });
 	}
 	return await pdfGetStructure(buf, password, dataProvider, {
+		nativeONNXRun,
 		sourceHash,
 	});
 }
@@ -59,6 +61,7 @@ async function getStructuredDocumentText(buf, options = {}) {
 		contentType: options.contentType,
 		password: options.password,
 		dataProvider: options.dataProvider,
+		nativeONNXRun: options.nativeONNXRun,
 		sourceHash: options.sourceHash,
 	});
 	let buffer = packStructuredDocumentText(structure, {
@@ -110,9 +113,9 @@ if (typeof self !== 'undefined') {
 	let waitingPromises = {};
 
 	self.query = async function (action, data, transfer) {
-		return new Promise(function (resolve) {
+		return new Promise(function (resolve, reject) {
 			promiseID++;
-			waitingPromises[promiseID] = resolve;
+			waitingPromises[promiseID] = { resolve, reject };
 			self.postMessage({ id: promiseID, action, data }, transfer);
 		});
 	};
@@ -121,9 +124,15 @@ if (typeof self !== 'undefined') {
 		let message = e.data;
 
 		if (message.responseID) {
-			let resolve = waitingPromises[message.responseID];
-			if (resolve) {
-				resolve(message.data);
+			let waiting = waitingPromises[message.responseID];
+			if (waiting) {
+				delete waitingPromises[message.responseID];
+				if (message.error) {
+					waiting.reject(message.error);
+				}
+				else {
+					waiting.resolve(message.data);
+				}
 			}
 			return;
 		}
@@ -274,6 +283,7 @@ if (typeof self !== 'undefined') {
 					contentType: message.data.contentType,
 					password: message.data.password,
 					dataProvider: fetchData,
+					nativeONNXRun: message.data.nativeONNX ? data => query('NativeONNXRun', data) : null,
 					sourceHash: message.data.sourceHash,
 				});
 				self.postMessage({

src/pdf/index.js

@@ -659,8 +659,18 @@ async function getStructure(buf, password, dataProvider, options = {}) {
 	let pdfManager = await getPdfManager(buf);
 	setHandler(pdfManager.pdfDocument, dataProvider);
 
-	let onnxRuntimeProvider = () => dataProvider('onnx/ort-wasm-simd.wasm');
-	let modelProvider = (name) => dataProvider(name.includes('/') ? name : name + '/model.onnx');
+	let useNativeONNX = typeof options.nativeONNXRun === 'function';
+	let onnxRuntimeProvider = useNativeONNX
+		? { type: 'native', run: options.nativeONNXRun }
+		: () => dataProvider('onnx/ort-wasm-simd.wasm');
+	let modelProvider = async (name) => {
+		let path = name.includes('/') ? name : name + '/model.onnx';
+		if (useNativeONNX && path.endsWith('.onnx')) {
+			return { nativeONNXModel: path };
+		}
+		let data = await dataProvider(path);
+		return data;
+	};
 	return await getFullStructure(pdfManager.pdfDocument, onnxRuntimeProvider, modelProvider, options);
 }
 

src/pdf/structure/model/onnx/host-onnx-runtime.js

@@ -0,0 +1,91 @@
+class NativeTensor {
+	constructor(type, data, dims) {
+		this.type = type;
+		this.data = data;
+		this.dims = dims;
+	}
+}
+
+class NativeInferenceSession {
+	constructor(model, nativeONNXRun) {
+		this.model = model;
+		this.nativeONNXRun = nativeONNXRun;
+	}
+
+	static async create(model, _options) {
+		const modelName = model?.nativeONNXModel;
+		const nativeONNXRun = model?.nativeONNXRun;
+		if (!modelName || typeof nativeONNXRun !== 'function') {
+			throw new Error('Native ONNX sessions require a native model provider');
+		}
+		return new NativeInferenceSession(modelName, nativeONNXRun);
+	}
+
+	async run(feeds, outputNames) {
+		const inputs = Object.entries(feeds).map(([name, tensor]) => ({
+			name,
+			type: tensor.type,
+			dims: Array.from(tensor.dims),
+			values: serializeTensorValues(tensor),
+		}));
+		const result = await this.nativeONNXRun({
+			model: this.model,
+			inputs,
+			outputNames: outputNames ?? null,
+		});
+		return deserializeOutputs(result?.outputs);
+	}
+
+	async release() {
+		// Native sessions are cached and owned by the host runtime.
+	}
+}
+
+export function createNativeRuntime(nativeONNXRun) {
+	return {
+		Tensor: NativeTensor,
+		InferenceSession: {
+			create: (model, options) => NativeInferenceSession.create({
+				...model,
+				nativeONNXRun,
+			}, options),
+		},
+	};
+}
+
+function serializeTensorValues(tensor) {
+	if (tensor.type === 'int64' && typeof BigInt64Array !== 'undefined' && tensor.data instanceof BigInt64Array) {
+		return Array.from(tensor.data, Number);
+	}
+	return Array.from(tensor.data);
+}
+
+function deserializeOutputs(outputs) {
+	if (!outputs || typeof outputs !== 'object') {
+		throw new Error('Native ONNX did not return outputs');
+	}
+	const result = {};
+	for (const [name, output] of Object.entries(outputs)) {
+		result[name] = new NativeTensor(output.type, typedArray(output.type, output.values), output.dims);
+	}
+	return result;
+}
+
+function typedArray(type, values) {
+	switch (type) {
+		case 'float32':
+			return Float32Array.from(values);
+
+		case 'int64':
+			if (typeof BigInt64Array === 'undefined') {
+				throw new Error('BigInt64Array is not available');
+			}
+			return BigInt64Array.from(values.map(BigInt));
+
+		case 'bool':
+			return Uint8Array.from(values);
+
+		default:
+			throw new Error(`Unsupported native ONNX output type: ${type}`);
+	}
+}

src/pdf/structure/model/onnx/runtime.js

@@ -1,9 +1,13 @@
 import * as ort from './ort.wasm.min.js';
 import { Mutex } from '../../../mutex.js';
+import { createNativeRuntime } from './host-onnx-runtime.js';
 
 export let onnxMutex = new Mutex()
 
 export async function getRuntime(onnxRuntimeProvider) {
+	if (onnxRuntimeProvider?.type === 'native') {
+		return createNativeRuntime(onnxRuntimeProvider.run);
+	}
 	return await onnxMutex.runExclusive(async () => {
 		ort.env.wasm.simd = true;
 		ort.env.wasm.numThreads = 1;