Use dynamic imports for onnxruntime-web

packages/web/src/onnx.ts

@@ -1,21 +1,32 @@
 export { createOnnxSession, runOnnxSession };
 
 import ndarray, { NdArray } from 'ndarray';
-import type ORT from 'onnxruntime-web';
-
-import * as ort_cpu from 'onnxruntime-web';
-import * as ort_gpu from 'onnxruntime-web/webgpu';
-
+import { InferenceSession, Tensor } from 'onnxruntime-web';
 import * as caps from './capabilities';
+import { loadAsUrl } from './resource';
 import { Config } from './schema';
-import { loadAsUrl, resolveChunkUrls } from './resource';
+
+type ORT = typeof import('onnxruntime-web');
+// use a dynamic import to avoid bundling the entire onnxruntime-web package
+let ort: ORT | null = null;
+const getOrt = async (useWebGPU: boolean): Promise<ORT> => {
+  if (ort !== null) {
+    return ort;
+  }
+  if (useWebGPU) {
+    ort = (await import('onnxruntime-web/webgpu')).default;
+  } else {
+    ort = (await import('onnxruntime-web')).default;
+  }
+  return ort;
+};
 
 async function createOnnxSession(model: any, config: Config) {
   const useWebGPU = config.device === 'gpu' && (await caps.webgpu());
   // BUG: proxyToWorker is not working for WASM/CPU Backend for now
   const proxyToWorker = useWebGPU && config.proxyToWorker;
   const executionProviders = [useWebGPU ? 'webgpu' : 'wasm'];
-  const ort = useWebGPU ? ort_gpu : ort_cpu;
+  const ort = await getOrt(useWebGPU);
 
   if (config.debug) {
     console.debug('\tUsing WebGPU:', useWebGPU);
@@ -45,14 +56,14 @@ async function createOnnxSession(model: any, config: Config) {
     console.debug('ort.env.wasm:', ort.env.wasm);
   }
 
-  const ort_config: ORT.InferenceSession.SessionOptions = {
+  const ortConfig: InferenceSession.SessionOptions = {
     executionProviders: executionProviders,
     graphOptimizationLevel: 'all',
     executionMode: 'parallel',
     enableCpuMemArena: true
   };
 
-  const session = await ort.InferenceSession.create(model, ort_config).catch(
+  const session = await ort.InferenceSession.create(model, ortConfig).catch(
     (e: any) => {
       throw new Error(
         `Failed to create session: "${e}". Please check if the publicPath is set correctly.`
@@ -69,7 +80,7 @@ async function runOnnxSession(
   config: Config
 ) {
   const useWebGPU = config.device === 'gpu' && (await caps.webgpu());
-  const ort = useWebGPU ? ort_gpu : ort_cpu;
+  const ort = await getOrt(useWebGPU);
 
   const feeds: Record<string, any> = {};
   for (const [key, tensor] of inputs) {
@@ -82,7 +93,7 @@ async function runOnnxSession(
   const outputData = await session.run(feeds, {});
   const outputKVPairs: NdArray<Float32Array>[] = [];
   for (const key of outputs) {
-    const output: ORT.Tensor = outputData[key];
+    const output: Tensor = outputData[key];
     const shape: number[] = output.dims as number[];
     const data: Float32Array = output.data as Float32Array;
     const tensor = ndarray(data, shape);