yizhi-js-lib/ai-box
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

增加MNN模型支持

Browse Source
This commit is contained in:
Yizhi
2025-03-07 17:18:20 +08:00
parent 4a6d092de1
commit bd90f2f6f6
24 changed files with 497 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/deploy/common/model.ts
View File

@ -45,6 +45,14 @@ export abstract class Model {
return new this(new backend.ort.Session(modelData as Uint8Array));
}
public static async fromMNN<T extends Model>(this: ModelConstructor<T>, modelData: Uint8Array | string) {
if (typeof modelData === "string") {
if (/^https?:\/\//.test(modelData)) modelData = await fetch(modelData).then(res => res.arrayBuffer()).then(buffer => new Uint8Array(buffer));
else modelData = await import("fs").then(fs => fs.promises.readFile(modelData as string));
}
return new this(new backend.mnn.Session(modelData as Uint8Array));
}
protected static async cacheModel<T extends Model, Create extends boolean = false>(this: ModelConstructor<T>, url: string, option?: ModelCacheOption<Create>): Promise<ModelCacheResult<T, Create>> {
//初始化目录
const [fs, path, os, crypto] = await Promise.all([import("fs"), import("path"), import("os"), import("crypto")]);
@ -111,6 +119,7 @@ export abstract class Model {
let model: T | undefined = undefined;
if (option?.createModel) {
if (modelType === "onnx") model = (this as any).fromOnnx(modelPath);
else if (modelType == "mnn") model = (this as any).fromMNN(modelPath);
}
return { modelPath, modelType, model: model as any }