Skip to content

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

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.

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

Merged
merged 5 commits into from
Apr 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 2 additions & 1 deletion src/env.js
Original file line number Diff line number Diff line change
Expand Up @@ -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} */
Expand Down
7 changes: 4 additions & 3 deletions src/models.js
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}));
}
Expand Down
16 changes: 11 additions & 5 deletions src/utils/hub.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.");
}

/**
Expand Down
Loading