feat: allow WASM modules to be imported via a user-land defined loade…

…r function microsoft#20876
kyr0 · Jul 20, 2024 · 1f968a8 · 1f968a8
1 parent 11bf309
commit 1f968a8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/js/common/lib/env.ts b/js/common/lib/env.ts
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 import {env as envImpl} from './env-impl.js';
+import type {OrtWasmModule} from 'onnxruntime-web';
 
 export declare namespace Env {
   export type WasmPathPrefix = string;
@@ -80,6 +81,15 @@ export declare namespace Env {
      * @defaultValue `false`
      */
     proxy?: boolean;
+
+    /**
+     * Set a custom function to import the WebAssembly module from user-land code (Inversion of Control, IoC).
+     *
+     * @defaultValue undefined
+     */
+    importWasmModule?: (mjsPathOverride: string, wasmPrefixOverride: string, numThreads: boolean) => Promise<[
+      string | undefined, EmscriptenModuleFactory<OrtWasmModule>
+    ]>;
   }
 
   export interface WebGLFlags {

diff --git a/js/web/lib/wasm/wasm-factory.ts b/js/web/lib/wasm/wasm-factory.ts
@@ -108,8 +108,9 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
   const mjsPathOverride = (mjsPathOverrideFlag as URL)?.href ?? mjsPathOverrideFlag;
   const wasmPathOverrideFlag = (wasmPaths as Env.WasmFilePaths)?.wasm;
   const wasmPathOverride = (wasmPathOverrideFlag as URL)?.href ?? wasmPathOverrideFlag;
+  const importFunction = typeof flags.importWasmModule === "function" ? flags.importWasmModule : importWasmModule;
 
-  const [objectUrl, ortWasmFactory] = (await importWasmModule(mjsPathOverride, wasmPrefixOverride, numThreads > 1));
+  const [objectUrl, ortWasmFactory] = (await importFunction(mjsPathOverride, wasmPrefixOverride, numThreads > 1));
 
   let isTimeout = false;