Skip to content

Commit 4cb6baf

Browse files
committed
added cross origin storage cache
1 parent f36e0bf commit 4cb6baf

File tree

2 files changed

+76
-0
lines changed

2 files changed

+76
-0
lines changed

src/utils/CrossOriginStorage.js

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
class CrossOriginStorage {
2+
static isAvailable = () => "crossOriginStorage" in navigator;
3+
4+
match = async (request) => {
5+
const hashValue = await this._getFileHash(request);
6+
if (!hashValue) {
7+
return undefined;
8+
}
9+
const hash = { algorithm: "SHA-256", value: hashValue };
10+
try {
11+
// @ts-expect-error
12+
const [handle] = await navigator.crossOriginStorage.requestFileHandles([
13+
hash,
14+
]);
15+
const blob = await handle.getFile();
16+
return new Response(blob);
17+
} catch (err) {
18+
return undefined;
19+
}
20+
};
21+
put = async (request, response) => {
22+
const blob = await response.blob();
23+
const hash = await this._getBlobHash(blob);
24+
// @ts-expect-error
25+
const [handle] = await navigator.crossOriginStorage.requestFileHandles(
26+
[hash],
27+
{ create: true },
28+
);
29+
const writableStream = await handle.createWritable();
30+
await writableStream.write(blob);
31+
await writableStream.close();
32+
};
33+
34+
_getFileHash = async (url) => {
35+
if (/\/resolve\/main\/onnx\//.test(url)) {
36+
const rawUrl = url.replace(/\/resolve\//, "/raw/");
37+
const text = await fetch(rawUrl).then((response) => response.text());
38+
if (!text.includes("oid sha256:")) {
39+
return null;
40+
}
41+
return text.replace(/.*?\n^oid sha256:(\w+)\n.*?$/gm, "$1") || null;
42+
}
43+
return null;
44+
};
45+
46+
_getBlobHash = async (blob) => {
47+
const hashAlgorithmIdentifier = "SHA-256";
48+
49+
// Get the contents of the blob as binary data contained in an ArrayBuffer.
50+
const arrayBuffer = await blob.arrayBuffer();
51+
52+
// Hash the arrayBuffer using SHA-256.
53+
const hashBuffer = await crypto.subtle.digest(
54+
hashAlgorithmIdentifier,
55+
arrayBuffer,
56+
);
57+
58+
// Convert the ArrayBuffer to a hex string.
59+
const hashArray = Array.from(new Uint8Array(hashBuffer));
60+
const hashHex = hashArray
61+
.map((byte) => byte.toString(16).padStart(2, "0"))
62+
.join("");
63+
64+
return {
65+
algorithm: hashAlgorithmIdentifier,
66+
value: hashHex,
67+
};
68+
};
69+
}
70+
71+
export default CrossOriginStorage;

src/utils/hub.js

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import path from 'node:path';
1010

1111
import { apis, env } from '../env.js';
1212
import { dispatchCallback } from './core.js';
13+
import CrossOriginStorage from './CrossOriginStorage.js'
1314

1415
/**
1516
* @typedef {boolean|number} ExternalData Whether to load the model using the external data format (used for models >= 2GB in size).
@@ -479,6 +480,10 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
479480
filename
480481
);
481482

483+
if(CrossOriginStorage.isAvailable()){
484+
cache = new CrossOriginStorage();
485+
}
486+
482487
/** @type {string} */
483488
let cacheKey;
484489
const proposedCacheKey = cache instanceof FileCache

0 commit comments

Comments
 (0)