Skip to content

Commit 7107711

Browse files
committed
stamp: add model loading fromStorage
- Adding `fromStorage` method to InferenceAPI, its allows model loadingfrom Supabase Storage with public/private bucket support.
1 parent eb1bb49 commit 7107711

File tree

2 files changed

+45
-2
lines changed

2 files changed

+45
-2
lines changed

ext/ai/js/onnxruntime/inference_api.js

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,13 @@ const DEFAULT_HUGGING_FACE_OPTIONS = {
1010
},
1111
};
1212

13+
const DEFAULT_STORAGE_OPTIONS = () => ({
14+
hostname: Deno.env.get('SUPABASE_URL'),
15+
mode: {
16+
authorization: Deno.env.get('SUPABASE_SERVICE_ROLE_KEY'),
17+
},
18+
});
19+
1320
/**
1421
* An user friendly API for onnx backend
1522
*/
@@ -28,14 +35,17 @@ class UserInferenceSession {
2835
this.outputs = session.outputNames;
2936
}
3037

31-
static async fromUrl(modelUrl) {
38+
static async fromUrl(modelUrl, authorization) {
3239
if (modelUrl instanceof URL) {
3340
modelUrl = modelUrl.toString();
3441
}
3542

3643
const encoder = new TextEncoder();
3744
const modelUrlBuffer = encoder.encode(modelUrl);
38-
const session = await InferenceSession.fromBuffer(modelUrlBuffer);
45+
const session = await InferenceSession.fromRequest(
46+
modelUrlBuffer,
47+
authorization,
48+
);
3949

4050
return new UserInferenceSession(session);
4151
}
@@ -61,6 +71,26 @@ class UserInferenceSession {
6171
return await UserInferenceSession.fromUrl(new URL(modelPath, hostname));
6272
}
6373

74+
static async fromStorage(modelPath, opts = {}) {
75+
const defaultOpts = DEFAULT_STORAGE_OPTIONS();
76+
const hostname = opts?.hostname ?? defaultOpts.hostname;
77+
const mode = opts?.mode ?? defaultOpts.mode;
78+
79+
const assetPath = mode === 'public' ? `public/${modelPath}` : `authenticated/${modelPath}`;
80+
81+
const storageUrl = `/storage/v1/object/${assetPath}`;
82+
83+
if (!URL.canParse(storageUrl, hostname)) {
84+
throw Error(
85+
`[Invalid URL] Couldn't parse the model path: "${storageUrl}"`,
86+
);
87+
}
88+
89+
return await UserInferenceSession.fromUrl(
90+
new URL(storageUrl, hostname),
91+
mode?.authorization,
92+
);
93+
}
6494
async run(inputs) {
6595
const outputs = await core.ops.op_ai_ort_run_session(this.id, inputs);
6696

types/global.d.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,19 @@ export namespace ai {
286286
};
287287
}): Promise<RawSession>;
288288

289+
/** Loads a ONNX model session from **Storage**.
290+
* Sessions are loaded once, then will keep warm cross worker's requests
291+
*/
292+
static fromStorage(repoId: string, opts?: {
293+
/**
294+
* @default 'env SUPABASE_URL'
295+
*/
296+
hostname?: string | URL;
297+
mode?: 'public' | {
298+
authorization: string;
299+
};
300+
}): Promise<RawSession>;
301+
289302
/** Run the current session with the given inputs.
290303
* Use `inputs` and `outputs` properties to know the required inputs and expected results for the model session.
291304
*

0 commit comments

Comments
 (0)