Add support for node.js external data & optimize downloading

Author: xenova

src/models.js

@@ -250,7 +250,7 @@ async function getSession(pretrained_model_name_or_path, fileName, options) {
 
     // handle onnx external data files
     const use_external_data_format = options.use_external_data_format ?? custom_config.use_external_data_format;
-    /** @type {Promise<{path: string, data: Uint8Array}>[]} */
+    /** @type {Promise<string|{path: string, data: Uint8Array}>[]} */
     let externalDataPromises = [];
     if (use_external_data_format) {
         let external_data_format;
@@ -272,8 +272,8 @@ async function getSession(pretrained_model_name_or_path, fileName, options) {
             const path = `${fileName}${suffix}.onnx_data${i === 0 ? '' : '_' + i}`;
             const fullPath = `${options.subfolder ?? ''}/${path}`;
             externalDataPromises.push(new Promise(async (resolve, reject) => {
-                const data = /** @type {Uint8Array} */ (await getModelFile(pretrained_model_name_or_path, fullPath, true, options, false));
-                resolve({ path, data })
+                const data = await getModelFile(pretrained_model_name_or_path, fullPath, true, options, apis.IS_NODE_ENV);
+                resolve(data instanceof Uint8Array ? { path, data } : path);
             }));
         }
 

src/utils/hub.js

@@ -8,7 +8,7 @@
 import fs from 'fs';
 import path from 'path';
 
-import { env } from '../env.js';
+import { apis, env } from '../env.js';
 import { dispatchCallback } from './core.js';
 
 /**
@@ -288,29 +288,61 @@ class FileCache {
     /**
      * Adds the given response to the cache.
      * @param {string} request 
-     * @param {Response|FileResponse} response 
+     * @param {Response} response 
+     * @param {(data: {progress: number, loaded: number, total: number}) => void} [progress_callback] Optional.
+     * The function to call with progress updates
      * @returns {Promise<void>}
      */
-    async put(request, response) {
-        const buffer = Buffer.from(await response.arrayBuffer());
-
-        let outputPath = path.join(this.path, request);
+    async put(request, response, progress_callback = undefined) {
+        let filePath = path.join(this.path, request);
 
         try {
-            await fs.promises.mkdir(path.dirname(outputPath), { recursive: true });
-            await fs.promises.writeFile(outputPath, buffer);
+            const contentLength = response.headers.get('Content-Length');
+            const total = parseInt(contentLength ?? '0');
+            let loaded = 0;
+
+            await fs.promises.mkdir(path.dirname(filePath), { recursive: true });
+            const fileStream = fs.createWriteStream(filePath);
+            const reader = response.body.getReader();
+
+            while (true) {
+                const { done, value } = await reader.read();
+                if (done) {
+                    break;
+                }
+
+                await new Promise((resolve, reject) => {
+                    fileStream.write(value, (err) => {
+                        if (err) {
+                            reject(err);
+                            return;
+                        }
+                        resolve();
+                    });
+                });
+
+                loaded += value.length;
+                const progress = total ? (loaded / total) * 100 : 0;
+
+                progress_callback?.({ progress, loaded, total });
+            }
 
-        } catch (err) {
-            console.warn('An error occurred while writing the file to cache:', err)
+            fileStream.close();
+        } catch (error) {
+            // Clean up the file if an error occurred during download
+            try {
+                await fs.promises.unlink(filePath);
+            } catch { }
+            throw error;
         }
-    }
 
-    // TODO add the rest?
-    // addAll(requests: RequestInfo[]): Promise<void>;
-    // delete(request: RequestInfo | URL, options?: CacheQueryOptions): Promise<boolean>;
-    // keys(request?: RequestInfo | URL, options?: CacheQueryOptions): Promise<ReadonlyArray<Request>>;
-    // match(request: RequestInfo | URL, options?: CacheQueryOptions): Promise<Response | undefined>;
-    // matchAll(request?: RequestInfo | URL, options?: CacheQueryOptions): Promise<ReadonlyArray<Response>>;
+        // TODO add the rest?
+        // addAll(requests: RequestInfo[]): Promise<void>;
+        // delete(request: RequestInfo | URL, options?: CacheQueryOptions): Promise<boolean>;
+        // keys(request?: RequestInfo | URL, options?: CacheQueryOptions): Promise<ReadonlyArray<Request>>;
+        // match(request: RequestInfo | URL, options?: CacheQueryOptions): Promise<Response | undefined>;
+        // matchAll(request?: RequestInfo | URL, options?: CacheQueryOptions): Promise<ReadonlyArray<Response>>;
+    }
 }
 
 /**
@@ -512,41 +544,45 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
         file: filename
     })
 
-    /** @type {Uint8Array} */
-    let buffer;
-
-    if (!options.progress_callback) {
-        // If no progress callback is specified, we can use the `.arrayBuffer()`
-        // method to read the response.
-        buffer = new Uint8Array(await response.arrayBuffer());
-
-    } else if (
-        cacheHit // The item is being read from the cache
-        &&
-        typeof navigator !== 'undefined' && /firefox/i.test(navigator.userAgent) // We are in Firefox
-    ) {
-        // Due to bug in Firefox, we cannot display progress when loading from cache.
-        // Fortunately, since this should be instantaneous, this should not impact users too much.
-        buffer = new Uint8Array(await response.arrayBuffer());
-
-        // For completeness, we still fire the final progress callback
-        dispatchCallback(options.progress_callback, {
-            status: 'progress',
-            name: path_or_repo_id,
-            file: filename,
-            progress: 100,
-            loaded: buffer.length,
-            total: buffer.length,
-        })
-    } else {
-        buffer = await readResponse(response, data => {
+    let result;
+    if (!(apis.IS_NODE_ENV && return_path)) {
+        /** @type {Uint8Array} */
+        let buffer;
+
+        if (!options.progress_callback) {
+            // If no progress callback is specified, we can use the `.arrayBuffer()`
+            // method to read the response.
+            buffer = new Uint8Array(await response.arrayBuffer());
+
+        } else if (
+            cacheHit // The item is being read from the cache
+            &&
+            typeof navigator !== 'undefined' && /firefox/i.test(navigator.userAgent) // We are in Firefox
+        ) {
+            // Due to bug in Firefox, we cannot display progress when loading from cache.
+            // Fortunately, since this should be instantaneous, this should not impact users too much.
+            buffer = new Uint8Array(await response.arrayBuffer());
+
+            // For completeness, we still fire the final progress callback
             dispatchCallback(options.progress_callback, {
                 status: 'progress',
                 name: path_or_repo_id,
                 file: filename,
-                ...data,
+                progress: 100,
+                loaded: buffer.length,
+                total: buffer.length,
             })
-        })
+        } else {
+            buffer = await readResponse(response, data => {
+                dispatchCallback(options.progress_callback, {
+                    status: 'progress',
+                    name: path_or_repo_id,
+                    file: filename,
+                    ...data,
+                })
+            })
+        }
+        result = buffer;
     }
 
     if (
@@ -557,36 +593,43 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
         // Check again whether request is in cache. If not, we add the response to the cache
         (await cache.match(cacheKey) === undefined)
     ) {
-        // NOTE: We use `new Response(buffer, ...)` instead of `response.clone()` to handle LFS files
-        await cache.put(cacheKey, new Response(buffer, {
-            headers: response.headers
-        }))
-            .catch(err => {
-                // Do not crash if unable to add to cache (e.g., QuotaExceededError).
-                // Rather, log a warning and proceed with execution.
-                console.warn(`Unable to add response to browser cache: ${err}.`);
-            });
+        if (!result) {
+            // We haven't yet read the response body, so we need to do so now.
+            await cache.put(cacheKey, /** @type {Response} */(response), options.progress_callback);
+        } else {
+            // NOTE: We use `new Response(buffer, ...)` instead of `response.clone()` to handle LFS files
+            await cache.put(cacheKey, new Response(result, {
+                headers: response.headers
+            }))
+                .catch(err => {
+                    // Do not crash if unable to add to cache (e.g., QuotaExceededError).
+                    // Rather, log a warning and proceed with execution.
+                    console.warn(`Unable to add response to browser cache: ${err}.`);
+                });
+        }
     }
     dispatchCallback(options.progress_callback, {
         status: 'done',
         name: path_or_repo_id,
         file: filename
     });
 
-    if (return_path) {
-        if (response instanceof FileResponse) {
-            return response.filePath;
-        } else {
-            const path = await cache.match(cacheKey);
-            if (path instanceof FileResponse) {
-                return path.filePath;
-            } else {
-                throw new Error("Unable to return path for response.");
-            }
+    if (result) {
+        if (return_path) {
+            throw new Error("Cannot return path in a browser environment.")
         }
+        return result;
+    }
+    if (response instanceof FileResponse) {
+        return response.filePath;
     }
 
-    return buffer;
+    const path = await cache.match(cacheKey);
+    if (path instanceof FileResponse) {
+        return path.filePath;
+    }
+    throw new Error("Unable to return path for response.");
+
 }
 
 /**