@aislamov/onnxruntime-web64
Version:
A Javascript library for running ONNX models on browsers
122 lines • 5.35 kB
JavaScript
;
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.OnnxruntimeWebAssemblySessionHandler = void 0;
const fs_1 = require("fs");
const onnxruntime_common_1 = require("onnxruntime-common");
const util_1 = require("util");
const proxy_wrapper_1 = require("./proxy-wrapper");
const wasm_common_1 = require("./wasm-common");
let runtimeInitialized;
class OnnxruntimeWebAssemblySessionHandler {
async fetchModelAndWeights(modelPath, weightsPath) {
const modelResponse = await fetch(modelPath);
if (modelResponse.status !== 200) {
throw new Error(`failed to load model: ${modelPath}`);
}
const promises = [
modelResponse.arrayBuffer().then(b => new Uint8Array(b))
];
if (weightsPath) {
const weightsResponse = await fetch(weightsPath);
const weightsSize = parseInt(weightsResponse.headers.get('Content-Length'), 10);
// we cannot create ArrayBuffer > 2gb but 64bit WASM Memory can have arbitrary size
const weightsMemory = new WebAssembly.Memory({
initial: Math.ceil(weightsSize / 65536),
maximum: Math.ceil(weightsSize / 65536),
// WASM Memory "index" parameter spec change landed but types are not yet updated
// https://github.com/WebAssembly/memory64/pull/39
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
index: 'u64',
shared: true,
});
promises.push((0, wasm_common_1.streamResponseToBuffer)(weightsResponse, weightsMemory.buffer, 0).then(() => weightsMemory.buffer));
}
// fetch model and weights in parallel
return Promise.all(promises);
}
async loadModel(urisOrBuffers, options) {
if (!runtimeInitialized) {
await (0, proxy_wrapper_1.initializeRuntime)(onnxruntime_common_1.env);
runtimeInitialized = true;
}
let modelBuffer;
let weightsBuffer;
const isNode = typeof process !== 'undefined' && process.versions && process.versions.node;
if (Array.isArray(urisOrBuffers)) {
// handle [string, string]
if (typeof urisOrBuffers[0] === 'string') {
if (isNode) {
modelBuffer = await (0, util_1.promisify)(fs_1.readFile)(urisOrBuffers[0]);
weightsBuffer = await (0, util_1.promisify)(fs_1.readFile)(urisOrBuffers[1]);
}
else {
[modelBuffer, weightsBuffer] = await this.fetchModelAndWeights(urisOrBuffers[0], urisOrBuffers[1]);
}
}
else { // [UInt8Array, ArrayBuffer]
[modelBuffer, weightsBuffer] = urisOrBuffers;
}
}
else {
if (typeof urisOrBuffers === 'string') {
if (isNode) {
modelBuffer = await (0, util_1.promisify)(fs_1.readFile)(urisOrBuffers);
}
else {
[modelBuffer] = await this.fetchModelAndWeights(urisOrBuffers);
}
}
else {
modelBuffer = urisOrBuffers;
}
}
const modelData = await (0, proxy_wrapper_1.createSessionAllocate)(modelBuffer, weightsBuffer);
// create the session
[this.sessionId, this.inputNames, this.outputNames] = await (0, proxy_wrapper_1.createSessionFinalize)(modelData, options);
}
async dispose() {
return (0, proxy_wrapper_1.releaseSession)(this.sessionId);
}
async run(feeds, fetches, options) {
const inputArray = [];
const inputIndices = [];
Object.entries(feeds).forEach(kvp => {
const name = kvp[0];
const tensor = kvp[1];
const index = this.inputNames.indexOf(name);
if (index === -1) {
throw new Error(`invalid input '${name}'`);
}
inputArray.push(tensor);
inputIndices.push(index);
});
const outputIndices = [];
Object.entries(fetches).forEach(kvp => {
const name = kvp[0];
// TODO: support pre-allocated output
const index = this.outputNames.indexOf(name);
if (index === -1) {
throw new Error(`invalid output '${name}'`);
}
outputIndices.push(index);
});
const outputs = await (0, proxy_wrapper_1.run)(this.sessionId, inputIndices, inputArray.map(t => [t.type, t.dims, t.data]), outputIndices, options);
const result = {};
for (let i = 0; i < outputs.length; i++) {
result[this.outputNames[outputIndices[i]]] =
new onnxruntime_common_1.Tensor(outputs[i][0], outputs[i][2], outputs[i][1].map(i => Number(i)));
}
return result;
}
startProfiling() {
// TODO: implement profiling
}
endProfiling() {
void (0, proxy_wrapper_1.endProfiling)(this.sessionId);
}
}
exports.OnnxruntimeWebAssemblySessionHandler = OnnxruntimeWebAssemblySessionHandler;
//# sourceMappingURL=session-handler.js.map