在浏览器上使用transformers.js运行(WebGPU)RMBG-1.4进行抠图(背景移除)
说明:
- 首次发表日期:2024-08-28
- 官方Github仓库地址: https://github.com/xenova/transformers.js/tree/main/examples/remove-background-client
准备
下载onnx模型文件: https://huggingface.co/briaai/RMBG-1.4/tree/main/onnx, 有3个onnx模型文件:
- model.onnx
- model_fp16.onnx
- model_quantized.onnx
Clone仓库并打开示例文件夹:
bash
git clone https://github.com/xenova/transformers.js.git
cd transformers.js/examples/remove-background-client
创建public/models/briaai/RMBG-1.4/onnx
文件夹,并将下载的模型文件放入。
运行tree public/
,可以看到:
public/
└── models
└── briaai
└── RMBG-1.4
└── onnx
├── model.onnx
├── model_fp16.onnx
└── model_quantized.onnx
安装项目依赖:
bash
npm install
推理
修改main.js文件中的env配置,使用本地模型:
js
env.allowLocalModels = true;
有很多可用的env配置,详见 https://huggingface.co/docs/transformers.js/api/env
以下是我测试过得配置(仅测试,不需要在main.js中修改):
js
env.remoteHost = 'https://hf-mirror.com';
env.allowRemoteModels = false;
// 默认会读取之前的缓存,如果之前的缓存是错的,需要清空浏览器缓存,或者使用该配置不使用浏览器缓存
env.useBrowserCache = false;
使用model_quantized.onnx
默认使用的模型文件名为model_quantized.onnx,不需要做任何修改。
使用model.onnx
修改main.js,添加quantized: false
到AutoModel.from_pretrained
:
js
const model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
// Do not require config.json to be present in the repository
config: { model_type: 'custom' },
quantized: false
// model_file_name: 'model_fp16',
// quantized: false
});
使用model_fp16.onnx
修改main.js,AutoModel.from_pretrained
部分如下:
js
const model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
// Do not require config.json to be present in the repository
config: { model_type: 'custom' },
model_file_name: 'model_fp16',
quantized: false
});
运行
bash
npm run dev
然后打开浏览器,等待状态从Loading model...变为Ready后,上传图片
添加日志查看使用的模型文件
修改node_modules/@xenova/transformers/src/utils/hub.js
中的getModelFile
函数,添加日志打印localPath:
js
let requestURL = pathJoin(path_or_repo_id, filename);
let localPath = pathJoin(env.localModelPath, requestURL);
console.log("localPath: ", localPath);
bash
npm run build
npm run preview
浏览器打开链接,F12显示开发者工具,可以看到日志:
localPath: /models/briaai/RMBG-1.4/onnx/model_fp16.onnx