Return buffer instead of file_path if cache unavailable for model loading (#1280)

* Return buffer if cache unavailable for model file loading

* Minor cleanup

* Add back check

* Actually, original approach was better

* Support custom cache returning response or path

---------

Co-authored-by: Joshua Lochner <[email protected]>

Author: PrafulB

src/env.js

@@ -118,7 +118,8 @@ const localModelPath = RUNNING_LOCALLY
  * @property {string} cacheDir The directory to use for caching files with the file system. By default, it is `./.cache`.
  * @property {boolean} useCustomCache Whether to use a custom cache system (defined by `customCache`), defaults to `false`.
  * @property {Object} customCache The custom cache to use. Defaults to `null`. Note: this must be an object which
- * 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
+ * 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.
+ * 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>`.
  */
 
 /** @type {TransformersEnvironment} */

src/models.js

@@ -116,7 +116,7 @@ import { RawImage } from './utils/image.js';
 import { dynamic_time_warping, max, medianFilter } from './utils/maths.js';
 import { EosTokenCriteria, MaxLengthCriteria, StoppingCriteriaList } from './generation/stopping_criteria.js';
 import { LogitsSampler } from './generation/logits_sampler.js';
-import { apis } from './env.js';
+import { apis, env } from './env.js';
 
 import { WhisperGenerationConfig } from './models/whisper/generation_whisper.js';
 import { whisper_language_to_code } from './models/whisper/common_whisper.js';
@@ -248,7 +248,8 @@ async function getSession(pretrained_model_name_or_path, fileName, options) {
         );
     }
 
-    const bufferOrPathPromise = getModelFile(pretrained_model_name_or_path, modelFileName, true, options, apis.IS_NODE_ENV);
+    const return_path = apis.IS_NODE_ENV && env.useFSCache;
+    const bufferOrPathPromise = getModelFile(pretrained_model_name_or_path, modelFileName, true, options, return_path);
 
     // handle onnx external data files
     const use_external_data_format = options.use_external_data_format ?? custom_config.use_external_data_format;
@@ -276,7 +277,7 @@ async function getSession(pretrained_model_name_or_path, fileName, options) {
             const path = `${baseName}_data${i === 0 ? '' : '_' + i}`;
             const fullPath = `${options.subfolder ?? ''}/${path}`;
             externalDataPromises.push(new Promise(async (resolve, reject) => {
-                const data = await getModelFile(pretrained_model_name_or_path, fullPath, true, options, apis.IS_NODE_ENV);
+                const data = await getModelFile(pretrained_model_name_or_path, fullPath, true, options, return_path);
                 resolve(data instanceof Uint8Array ? { path, data } : path);
             }));
         }

src/utils/hub.js

@@ -638,7 +638,7 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
     });
 
     if (result) {
-        if (return_path) {
+        if (!apis.IS_NODE_ENV && return_path) {
             throw new Error("Cannot return path in a browser environment.")
         }
         return result;
@@ -647,12 +647,18 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
         return response.filePath;
     }
 
-    const path = await cache.match(cacheKey);
-    if (path instanceof FileResponse) {
-        return path.filePath;
+    // Otherwise, return the cached response (most likely a `FileResponse`).
+    // NOTE: A custom cache may return a Response, or a string (file path)
+    const cachedResponse = await cache?.match(cacheKey);
+    if (cachedResponse instanceof FileResponse) {
+        return cachedResponse.filePath;
+    } else if (cachedResponse instanceof Response) {
+        return new Uint8Array(await cachedResponse.arrayBuffer());
+    } else if (typeof cachedResponse === 'string') {
+        return cachedResponse;
     }
-    throw new Error("Unable to return path for response.");
 
+    throw new Error("Unable to get model file path or buffer.");
 }
 
 /**