Spaces:
Running
Running
| // Resumable, chunked model download into a keyed store. | |
| // | |
| // Pulled out of the worker so it can be tested in node with an in-memory store | |
| // (see test/download-test.mjs) - the chunking, Range and restart logic is the | |
| // part that is easy to get wrong. | |
| // | |
| // Store interface (default: storage.js / IndexedDB): | |
| // CHUNK_SIZE | |
| // getMeta(id) -> record or undefined | |
| // putMeta(record) -> full replace (in-memory test store) | |
| // updateMeta(id, patch) -> merge (optional; preferred when present) | |
| // putChunk(id, index, buffer) | |
| // clearChunks(id) | |
| // | |
| // The record carries the model's identity (id, name, source), so this module | |
| // only ever touches the fields it owns: total, received and complete. | |
| import { | |
| CHUNK_SIZE, getModel, putModel, updateModel, putChunk, clearChunks, | |
| } from "./storage.js"; | |
| const defaultStorage = { | |
| CHUNK_SIZE, getMeta: getModel, putMeta: putModel, updateMeta: updateModel, putChunk, clearChunks, | |
| }; | |
| /** | |
| * Ensure the whole file is stored, resuming a previous partial download. | |
| * | |
| * @param {{ id: string, url: string, total?: number }} target | |
| * @param {{ signal?: AbortSignal, onProgress?: Function, storage?: object }} options | |
| * @returns {Promise<object>} the stored model record | |
| */ | |
| export async function downloadModel(target, options = {}) { | |
| const { id, url } = target; | |
| const storage = options.storage || defaultStorage; | |
| const onProgress = options.onProgress || (() => {}); | |
| const signal = options.signal; | |
| const chunkSize = storage.CHUNK_SIZE; | |
| let meta = await storage.getMeta(id); | |
| if (meta && meta.complete) { | |
| onProgress({ phase: "download", loaded: meta.received, total: meta.total, complete: true }); | |
| return meta; | |
| } | |
| const writeMeta = async (fields) => { | |
| if (storage.updateMeta) return storage.updateMeta(id, fields); | |
| return storage.putMeta({ ...(meta || {}), id, url, ...fields }); | |
| }; | |
| let received = meta ? (meta.received || 0) : 0; | |
| let total = (meta && meta.total) || target.total || 0; | |
| // a partial chunk would not line up with the Range offset - start over | |
| if (received % chunkSize !== 0) { | |
| await storage.clearChunks(id); | |
| received = 0; | |
| } | |
| const headers = received > 0 ? { Range: `bytes=${received}-` } : undefined; | |
| const res = await fetch(url, { headers, signal }); | |
| if (!res.ok) { | |
| throw new Error(`model download failed: HTTP ${res.status} ${res.statusText}`); | |
| } | |
| // server ignored the Range header and sent the whole file: discard the partial copy | |
| if (received > 0 && res.status !== 206) { | |
| await storage.clearChunks(id); | |
| received = 0; | |
| } | |
| if (!total) { | |
| total = received + Number(res.headers.get("content-length") || 0); | |
| } | |
| let index = Math.floor(received / chunkSize); | |
| await writeMeta({ total, received, complete: false }); | |
| onProgress({ phase: "download", loaded: received, total }); | |
| const reader = res.body.getReader(); | |
| let pending = new Uint8Array(chunkSize); | |
| let fill = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| let src = value; | |
| while (src.length > 0) { | |
| const take = Math.min(src.length, chunkSize - fill); | |
| pending.set(src.subarray(0, take), fill); | |
| fill += take; | |
| src = src.subarray(take); | |
| if (fill === chunkSize) { | |
| await storage.putChunk(id, index++, pending.buffer); | |
| received += fill; | |
| fill = 0; | |
| pending = new Uint8Array(chunkSize); | |
| await writeMeta({ total, received, complete: false }); | |
| onProgress({ phase: "download", loaded: received, total }); | |
| } | |
| } | |
| } | |
| if (fill > 0) { | |
| await storage.putChunk(id, index++, pending.buffer.slice(0, fill)); | |
| received += fill; | |
| } | |
| await writeMeta({ total: total || received, received, complete: true }); | |
| onProgress({ phase: "download", loaded: received, total: total || received, complete: true }); | |
| return storage.getMeta(id); | |
| } | |