Skip to content

Commit 0082c20

Browse files
nico-martinxenova
nico-martin
and
xenova
authored
[v4] Cache wasm files to enable fully offline usage after initial load (#1471)
* added wasm cache * some refactoring of the hub.js and caching of the wasm factory * fixed comment * added string as cache return * fixes after review * Only return if match is found * Return response even if cache doesn't exist Don't throw error if we can't open cache or load file from cache, but we are able to make the request. --------- Co-authored-by: Joshua Lochner <[email protected]> Co-authored-by: Joshua Lochner <[email protected]>
1 parent aab2326 commit 0082c20

File tree

9 files changed

+660
-453
lines changed

9 files changed

+660
-453
lines changed

‎src/backends/onnx.js‎

Lines changed: 81 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import { env, apis } from '../env.js';
2222
// In either case, we select the default export if it exists, otherwise we use the named export.
2323
import * as ONNX_NODE from 'onnxruntime-node';
2424
import * as ONNX_WEB from 'onnxruntime-web/webgpu';
25+
import { loadWasmBinary, loadWasmFactory } from './utils/cacheWasm.js';
2526

2627
export { Tensor } from 'onnxruntime-common';
2728

@@ -145,6 +146,79 @@ const IS_WEB_ENV = apis.IS_BROWSER_ENV || apis.IS_WEBWORKER_ENV;
145146
*/
146147
let webInitChain = Promise.resolve();
147148

149+
/**
150+
* Promise that resolves when WASM binary has been loaded (if caching is enabled).
151+
* This ensures we only attempt to load the WASM binary once.
152+
* @type {Promise<void>|null}
153+
*/
154+
let wasmLoadPromise = null;
155+
156+
/**
157+
* Ensures the WASM binary is loaded and cached before creating an inference session.
158+
* Only runs once, even if called multiple times.
159+
*
160+
* @returns {Promise<void>}
161+
*/
162+
async function ensureWasmLoaded() {
163+
// If already loading or loaded, return the existing promise
164+
if (wasmLoadPromise) {
165+
return wasmLoadPromise;
166+
}
167+
168+
const shouldUseWasmCache =
169+
env.useWasmCache &&
170+
typeof ONNX_ENV?.wasm?.wasmPaths === 'object' &&
171+
ONNX_ENV?.wasm?.wasmPaths?.wasm &&
172+
ONNX_ENV?.wasm?.wasmPaths?.mjs;
173+
174+
// Check if we should load the WASM binary
175+
if (!shouldUseWasmCache) {
176+
wasmLoadPromise = Promise.resolve();
177+
return wasmLoadPromise;
178+
}
179+
180+
// Start loading the WASM binary
181+
wasmLoadPromise = (async () => {
182+
// At this point, we know wasmPaths is an object (not a string) because
183+
// shouldUseWasmCache checks for wasmPaths.wasm and wasmPaths.mjs
184+
const urls = /** @type {{ wasm: string, mjs: string }} */ (ONNX_ENV.wasm.wasmPaths);
185+
186+
// Load and cache both the WASM binary and factory
187+
await Promise.all([
188+
// Load and cache the WASM binary
189+
urls.wasm
190+
? (async () => {
191+
try {
192+
const wasmBinary = await loadWasmBinary(urls.wasm);
193+
if (wasmBinary) {
194+
ONNX_ENV.wasm.wasmBinary = wasmBinary;
195+
}
196+
} catch (err) {
197+
console.warn('Failed to pre-load WASM binary:', err);
198+
}
199+
})()
200+
: Promise.resolve(),
201+
202+
// Load and cache the WASM factory
203+
urls.mjs
204+
? (async () => {
205+
try {
206+
const wasmFactoryBlob = await loadWasmFactory(urls.mjs);
207+
if (wasmFactoryBlob) {
208+
// @ts-ignore
209+
ONNX_ENV.wasm.wasmPaths.mjs = wasmFactoryBlob;
210+
}
211+
} catch (err) {
212+
console.warn('Failed to pre-load WASM factory:', err);
213+
}
214+
})()
215+
: Promise.resolve(),
216+
]);
217+
})();
218+
219+
return wasmLoadPromise;
220+
}
221+
148222
/**
149223
* Create an ONNX inference session.
150224
* @param {Uint8Array|string} buffer_or_path The ONNX model buffer or path.
@@ -153,6 +227,7 @@ let webInitChain = Promise.resolve();
153227
* @returns {Promise<import('onnxruntime-common').InferenceSession & { config: Object}>} The ONNX inference session.
154228
*/
155229
export async function createInferenceSession(buffer_or_path, session_options, session_config) {
230+
await ensureWasmLoaded();
156231
const load = () => InferenceSession.create(buffer_or_path, {
157232
// Set default log level, but allow overriding through session options
158233
logSeverityLevel: DEFAULT_LOG_LEVEL,
@@ -209,15 +284,15 @@ if (ONNX_ENV?.wasm) {
209284

210285
ONNX_ENV.wasm.wasmPaths = apis.IS_SAFARI
211286
? {
212-
mjs: `${wasmPathPrefix}/ort-wasm-simd-threaded.mjs`,
213-
wasm: `${wasmPathPrefix}/ort-wasm-simd-threaded.wasm`,
287+
mjs: `${wasmPathPrefix}ort-wasm-simd-threaded.mjs`,
288+
wasm: `${wasmPathPrefix}ort-wasm-simd-threaded.wasm`,
214289
}
215-
: wasmPathPrefix;
290+
: {
291+
mjs: `${wasmPathPrefix}ort-wasm-simd-threaded.asyncify.mjs`,
292+
wasm: `${wasmPathPrefix}ort-wasm-simd-threaded.asyncify.wasm`,
293+
};
216294
}
217295

218-
// TODO: Add support for loading WASM files from cached buffer when we upgrade to [email protected]
219-
// https://github.com/microsoft/onnxruntime/pull/21534
220-
221296
// Users may wish to proxy the WASM backend to prevent the UI from freezing,
222297
// However, this is not necessary when using WebGPU, so we default to false.
223298
ONNX_ENV.wasm.proxy = false;

‎src/backends/utils/cacheWasm.js‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { getCache } from '../../utils/cache.js';
2+
3+
/**
4+
* Loads and caches a file from the given URL.
5+
* @param {string} url The URL of the file to load.
6+
* @returns {Promise<Response|import('../../utils/hub/FileResponse.js').default|null|string>} The response object, or null if loading failed.
7+
*/
8+
async function loadAndCacheFile(url) {
9+
const fileName = url.split('/').pop();
10+
11+
/** @type {import('../../utils/cache.js').CacheInterface|undefined} */
12+
let cache;
13+
try {
14+
cache = await getCache();
15+
16+
// Try to get from cache first
17+
if (cache) {
18+
const result = await cache.match(url);
19+
if (result) {
20+
return result;
21+
}
22+
}
23+
} catch (error) {
24+
console.warn(`Failed to load ${fileName} from cache:`, error);
25+
}
26+
27+
// If not in cache, fetch it
28+
const response = await fetch(url);
29+
30+
if (!response.ok) {
31+
throw new Error(`Failed to fetch ${fileName}: ${response.status} ${response.statusText}`);
32+
}
33+
34+
// Cache the response for future use
35+
if (cache) {
36+
try {
37+
await cache.put(url, response.clone());
38+
} catch (e) {
39+
console.warn(`Failed to cache ${fileName}:`, e);
40+
}
41+
}
42+
43+
return response;
44+
45+
}
46+
47+
/**
48+
* Loads and caches the WASM binary for ONNX Runtime.
49+
* @param {string} wasmURL The URL of the WASM file to load.
50+
* @returns {Promise<ArrayBuffer|null>} The WASM binary as an ArrayBuffer, or null if loading failed.
51+
*/
52+
53+
export async function loadWasmBinary(wasmURL) {
54+
const response = await loadAndCacheFile(wasmURL);
55+
if (!response || typeof response === 'string') return null;
56+
57+
try {
58+
return await response.arrayBuffer();
59+
} catch (error) {
60+
console.warn('Failed to read WASM binary:', error);
61+
return null;
62+
}
63+
}
64+
65+
/**
66+
* Loads and caches the WASM Factory for ONNX Runtime.
67+
* @param {string} libURL The URL of the WASM Factory to load.
68+
* @returns {Promise<string|null>} The blob URL of the WASM Factory, or null if loading failed.
69+
*/
70+
export async function loadWasmFactory(libURL) {
71+
const response = await loadAndCacheFile(libURL);
72+
if (!response || typeof response === 'string') return null;
73+
74+
try {
75+
let code = await response.text();
76+
// Fix relative paths when loading factory from blob, overwrite import.meta.url with actual baseURL
77+
const baseUrl = libURL.split('/').slice(0, -1).join('/');
78+
code = code.replace(/import\.meta\.url/g, `"${baseUrl}"`);
79+
const blob = new Blob([code], { type: 'text/javascript' });
80+
return URL.createObjectURL(blob);
81+
} catch (error) {
82+
console.warn('Failed to read WASM binary:', error);
83+
return null;
84+
}
85+
}

‎src/env.js‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,12 @@ const localModelPath = RUNNING_LOCALLY ? path.join(dirname__, DEFAULT_LOCAL_MODE
152152
* @property {boolean} useFSCache Whether to use the file system to cache files. By default, it is `true` if available.
153153
* @property {string|null} cacheDir The directory to use for caching files with the file system. By default, it is `./.cache`.
154154
* @property {boolean} useCustomCache Whether to use a custom cache system (defined by `customCache`), defaults to `false`.
155-
* @property {Object|null} customCache The custom cache to use. Defaults to `null`. Note: this must be an object which
155+
* @property {import('./utils/cache.js').CacheInterface|null} customCache The custom cache to use. Defaults to `null`. Note: this must be an object which
156156
* implements the `match` and `put` functions of the Web Cache API. For more information, see https://developer.mozilla.org/en-US/docs/Web/API/Cache.
157-
* If you wish, you may also return a `Promise<string>` from the `match` function if you'd like to use a file path instead of `Promise<Response>`.
157+
* @property {boolean} useWasmCache Whether to pre-load and cache WASM binaries for ONNX Runtime. Defaults to `true` when cache is available.
158+
* This can improve performance by avoiding repeated downloads of WASM files. Note: Only the WASM binary is cached.
159+
* The MJS loader file still requires network access unless you use a Service Worker.
160+
* @property {string} cacheKey The cache key to use for storing models and WASM binaries. Defaults to 'transformers-cache'.
158161
*/
159162

160163
/** @type {TransformersEnvironment} */
@@ -185,6 +188,9 @@ export const env = {
185188

186189
useCustomCache: false,
187190
customCache: null,
191+
192+
useWasmCache: IS_WEB_CACHE_AVAILABLE || IS_FS_AVAILABLE,
193+
cacheKey: 'transformers-cache',
188194
//////////////////////////////////////////////////////
189195
};
190196

‎src/utils/cache.js‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import { apis, env } from '../env.js';
2+
import FileCache from './hub/FileCache.js';
3+
4+
/**
5+
* @typedef {Object} CacheInterface
6+
* @property {(request: string) => Promise<Response|import('./hub/FileResponse.js').default|undefined|string>} match
7+
* Checks if a request is in the cache and returns the cached response if found.
8+
* @property {(request: string, response: Response, progress_callback?: (data: {progress: number, loaded: number, total: number}) => void) => Promise<void>} put
9+
* Adds a response to the cache.
10+
*/
11+
12+
/**
13+
* Retrieves an appropriate caching backend based on the environment configuration.
14+
* Attempts to use custom cache, browser cache, or file system cache in that order of priority.
15+
* @returns {Promise<CacheInterface | null>}
16+
* @param file_cache_dir {string|null} Path to a directory in which a downloaded pretrained model configuration should be cached if using the file system cache.
17+
*/
18+
export async function getCache(file_cache_dir = null) {
19+
// First, check if the a caching backend is available
20+
// If no caching mechanism available, will download the file every time
21+
let cache = null;
22+
if (env.useCustomCache) {
23+
// Allow the user to specify a custom cache system.
24+
if (!env.customCache) {
25+
throw Error('`env.useCustomCache=true`, but `env.customCache` is not defined.');
26+
}
27+
28+
// Check that the required methods are defined:
29+
if (!env.customCache.match || !env.customCache.put) {
30+
throw new Error(
31+
'`env.customCache` must be an object which implements the `match` and `put` functions of the Web Cache API. ' +
32+
'For more information, see https://developer.mozilla.org/en-US/docs/Web/API/Cache',
33+
);
34+
}
35+
cache = env.customCache;
36+
}
37+
38+
if (!cache && env.useBrowserCache) {
39+
if (typeof caches === 'undefined') {
40+
throw Error('Browser cache is not available in this environment.');
41+
}
42+
try {
43+
// In some cases, the browser cache may be visible, but not accessible due to security restrictions.
44+
// For example, when running an application in an iframe, if a user attempts to load the page in
45+
// incognito mode, the following error is thrown: `DOMException: Failed to execute 'open' on 'CacheStorage':
46+
// An attempt was made to break through the security policy of the user agent.`
47+
// So, instead of crashing, we just ignore the error and continue without using the cache.
48+
cache = await caches.open(env.cacheKey);
49+
} catch (e) {
50+
console.warn('An error occurred while opening the browser cache:', e);
51+
}
52+
}
53+
54+
if (!cache && env.useFSCache) {
55+
if (!apis.IS_FS_AVAILABLE) {
56+
throw Error('File System Cache is not available in this environment.');
57+
}
58+
59+
// If `cache_dir` is not specified, use the default cache directory
60+
cache = new FileCache(file_cache_dir ?? env.cacheDir);
61+
}
62+
63+
return cache;
64+
}
65+
66+
/**
67+
* Searches the cache for any of the provided names and returns the first match found.
68+
* @param {CacheInterface} cache The cache to search
69+
* @param {...string} names The names of the items to search for
70+
* @returns {Promise<import('./hub/FileResponse.js').default|Response|undefined|string>} The item from the cache, or undefined if not found.
71+
*/
72+
export async function tryCache(cache, ...names) {
73+
for (let name of names) {
74+
try {
75+
let result = await cache.match(name);
76+
if (result) return result;
77+
} catch (e) {
78+
continue;
79+
}
80+
}
81+
return undefined;
82+
}

0 commit comments

Comments
 (0)