Skip to content

[WIP] Deno browser WASM/WebGPU support #1344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[WIP] Deno browser WASM/WebGPU support #1344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions src/backends/onnx.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,12 +56,15 @@ let defaultDevices;
let ONNX;
const ORT_SYMBOL = Symbol.for('onnxruntime');

/** @type {"custom"|"node"|"web"} */
let ort;
if (ORT_SYMBOL in globalThis) {
// If the JS runtime exposes their own ONNX runtime, use it
ONNX = globalThis[ORT_SYMBOL];
ort = 'custom';

} else if (apis.IS_NODE_ENV) {
ONNX = ONNX_NODE.default ?? ONNX_NODE;
} else if (apis.IS_NODE_ENV && (ONNX = ONNX_NODE.default ?? ONNX_NODE)?.InferenceSession) {
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handles an edge case where it is a node-like environment (e.g., Deno), but we are running in an environment where onnxruntime-node is not bundled (but onnxruntime-wasm is)

ort = 'node';

// Updated as of ONNX Runtime 1.20.1
// The following table lists the supported versions of ONNX Runtime Node.js binding provided with pre-built binaries.
Expand All @@ -87,6 +90,7 @@ if (ORT_SYMBOL in globalThis) {
defaultDevices = ['cpu'];
} else {
ONNX = ONNX_WEB;
ort = 'web';

if (apis.IS_WEBNN_AVAILABLE) {
// TODO: Only push supported providers (depending on available hardware)
Expand Down Expand Up @@ -169,6 +173,14 @@ export function isONNXTensor(x) {
return x instanceof ONNX.Tensor;
}

/**
* The type of ONNX runtime being used.
* - 'node' for `onnxruntime-node`
* - 'web' for `onnxruntime-web`
* - 'custom' for a custom ONNX runtime
*/
export const runtime = ort;

/** @type {import('onnxruntime-common').Env} */
// @ts-ignore
const ONNX_ENV = ONNX?.env;
Expand Down
2 changes: 1 addition & 1 deletion src/env.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ export const env = {
remoteHost: 'https://huggingface.co/',
remotePathTemplate: '{model}/resolve/{revision}/',

allowLocalModels: !(IS_BROWSER_ENV || IS_WEBWORKER_ENV),
allowLocalModels: !(IS_BROWSER_ENV || IS_WEBWORKER_ENV || IS_DENO_RUNTIME),
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deno doesn't support relative URLs, so this is necessary... even if we have filesystem access

localModelPath: localModelPath,
useFS: IS_FS_AVAILABLE,

Expand Down
3 changes: 2 additions & 1 deletion src/models.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ import {
createInferenceSession,
isONNXTensor,
isONNXProxy,
runtime,
} from './backends/onnx.js';
import {
DATA_TYPES,
Expand Down Expand Up @@ -172,7 +173,7 @@ async function getSession(pretrained_model_name_or_path, fileName, options) {

// If the device is not specified, we use the default (supported) execution providers.
const selectedDevice = /** @type {import("./utils/devices.js").DeviceType} */(
device ?? (apis.IS_NODE_ENV ? 'cpu' : 'wasm')
device ?? (runtime === "web" ? 'wasm' : 'cpu')
);

const executionProviders = deviceToExecutionProviders(selectedDevice);
Expand Down
16 changes: 15 additions & 1 deletion webpack.config.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,21 @@ class PostBuildPlugin {
{
const src = path.join(__dirname, 'node_modules/onnxruntime-web/dist', ORT_JSEP_FILE);
const dest = path.join(dist, ORT_JSEP_FILE);
fs.copyFileSync(src, dest);

// Transformers.js uses both onnxruntime-web and onnxruntime-node in the same package,
// and the runtime we use depends on the environment (onnxruntime-web for web, onnxruntime-node for Node.js).
// This means that we don't currently support using the WASM backend in Node.js, so we disable this behaviour in the JSEP file.
const content = fs.readFileSync(src, 'utf8');
const updatedContent = content
.replace(
`"object"==typeof process&&"object"==typeof process.versions&&"string"==typeof process.versions.node&&"renderer"!=process.type`,
"false",
)
.replace(
`typeof globalThis.process?.versions?.node == 'string'`,
"false",
)
Comment on lines +63 to +70
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very hacky, but is a current workaround, otherwise we get errors related to worker_threads not being prefixed with node:

cc @fs-eire if you have a better idea.

fs.writeFileSync(dest, updatedContent);
}
});
}
Expand Down