Skip to content

Commit

Permalink
Add missing 'ready' status in the ProgressInfo type (#1070)
Browse files Browse the repository at this point in the history
  • Loading branch information
ocavue authored Dec 7, 2024
1 parent c850083 commit 7dffb9a
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 21 deletions.
48 changes: 39 additions & 9 deletions src/utils/core.js
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,45 @@
*/

/**
* @typedef {Object} ProgressInfo
* @property {'initiate' | 'download' | 'progress' | 'done'} status The status of the progress item.
* @property {string} name This can be either:
* - a string, the *model id* of a model repo on huggingface.co.
* - a path to a *directory* potentially containing the file.
* @property {string} file The name of the file
* @property {number} [progress] A number between 0 and 100. Only available for the 'progress' status.
* @property {number} [loaded] The number of bytes loaded. Only available for the 'progress' status.
* @property {number} [total] The total number of bytes to be loaded. Only available for the 'progress' status.
* @typedef {Object} InitiateProgressInfo
* @property {'initiate'} status
* @property {string} name The model id or directory path.
* @property {string} file The name of the file.
*/

/**
* @typedef {Object} DownloadProgressInfo
* @property {'download'} status
* @property {string} name The model id or directory path.
* @property {string} file The name of the file.
*/

/**
* @typedef {Object} ProgressStatusInfo
* @property {'progress'} status
* @property {string} name The model id or directory path.
* @property {string} file The name of the file.
* @property {number} progress A number between 0 and 100.
* @property {number} loaded The number of bytes loaded.
* @property {number} total The total number of bytes to be loaded.
*/

/**
* @typedef {Object} DoneProgressInfo
* @property {'done'} status
* @property {string} name The model id or directory path.
* @property {string} file The name of the file.
*/

/**
* @typedef {Object} ReadyProgressInfo
* @property {'ready'} status
* @property {string} task The loaded task.
* @property {string} model The loaded model.
*/

/**
* @typedef {InitiateProgressInfo | DownloadProgressInfo | ProgressStatusInfo | DoneProgressInfo | ReadyProgressInfo} ProgressInfo
*/

/**
Expand Down
20 changes: 8 additions & 12 deletions src/utils/hub.js
Original file line number Diff line number Diff line change
Expand Up @@ -504,13 +504,6 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
file: filename
})

/** @type {import('./core.js').ProgressInfo} */
const progressInfo = {
status: 'progress',
name: path_or_repo_id,
file: filename
}

/** @type {Uint8Array} */
let buffer;

Expand All @@ -530,15 +523,19 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti

// For completeness, we still fire the final progress callback
dispatchCallback(options.progress_callback, {
...progressInfo,
status: 'progress',
name: path_or_repo_id,
file: filename,
progress: 100,
loaded: buffer.length,
total: buffer.length,
})
} else {
buffer = await readResponse(response, data => {
dispatchCallback(options.progress_callback, {
...progressInfo,
status: 'progress',
name: path_or_repo_id,
file: filename,
...data,
})
})
Expand Down Expand Up @@ -595,12 +592,11 @@ export async function getModelJSON(modelPath, fileName, fatal = true, options =

return JSON.parse(jsonData);
}

/**
* Read and track progress when reading a Response object
*
* @param {any} response The Response object to read
* @param {function} progress_callback The function to call with progress updates
* @param {Response|FileResponse} response The Response object to read
* @param {(data: {progress: number, loaded: number, total: number}) => void} progress_callback The function to call with progress updates
* @returns {Promise<Uint8Array>} A Promise that resolves with the Uint8Array buffer
*/
async function readResponse(response, progress_callback) {
Expand Down

0 comments on commit 7dffb9a

Please sign in to comment.