happyme531/wd-convnext-tagger-v3-RKNN2

WD ConvNext Tagger v3 RKNN2

(English README see below)

在RK3588上运行WaifuDiffusion图像标签模型!

• 推理速度(RK3588):

  • 单NPU核: 320ms

• 内存占用(RK3588):

  • 0.45GB

使用方法

1. 克隆或者下载此仓库到本地

2. 安装依赖

pip install numpy<2 pandas opencv-python rknn-toolkit-lite2

3. 运行

python run_rknn.py input.jpg

输出结果示例:

       tag_id           name     probs
0     9999999        general  0.521484
5      212816           solo  0.929199
12      15080     short_hair  0.520508
25     540830           1boy  0.947754
40      16613        jewelry  0.577148
72    1300281     male_focus  0.907227
130     10926          pants  0.803223
346   1094664   colored_skin  0.570312
373      4009     turtleneck  0.552246
1532  1314823  black_sweater  0.514160

模型转换

1. 安装依赖

pip install numpy<2 onnxruntime rknn-toolkit2

2. 下载原始onnx模型

3. 转换onnx模型到rknn模型:

python convert_rknn.py

已知问题

• int8量化后精度损失极大, 基本不可用. 不建议使用量化推理.

参考

• SmilingWolf/wd-convnext-tagger-v3

English README

Run WaifuDiffusion image tagging model on RK3588!

• Inference Speed (RK3588):

  • Single NPU Core: 320ms

• Memory Usage (RK3588):

  • 0.45GB

Usage

1. Clone or download this repository

2. Install dependencies

pip install numpy<2 pandas opencv-python rknn-toolkit-lite2

3. Run

python run_rknn.py input.jpg

Output example:

       tag_id           name     probs
0     9999999        general  0.521484
5      212816           solo  0.929199
12      15080     short_hair  0.520508
25     540830           1boy  0.947754
40      16613        jewelry  0.577148
72    1300281     male_focus  0.907227
130     10926          pants  0.803223
346   1094664   colored_skin  0.570312
373      4009     turtleneck  0.552246
1532  1314823  black_sweater  0.514160

Model Conversion

1. Install dependencies

pip install numpy<2 onnxruntime rknn-toolkit2

2. Download original onnx model

3. Convert onnx model to rknn model:

python convert_rknn.py

Known Issues

• Huge precision loss after int8 quantization, not recommended to use quantized inference.

References

• SmilingWolf/wd-convnext-tagger-v3