pi-ux-ce/wd14-tagger-standalone
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support gpu.

Browse Source
This commit is contained in:
corkborg
2024-04-23 01:48:39 +09:00
parent 282472f212
commit babad14542
4 changed files with 37 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+19 -2
View File
@@ -13,7 +13,8 @@ pip install -r requirements.txt
## usage
```
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--gpu]
[--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
options:
-h, --help show this help message and exit
@@ -22,7 +23,8 @@ options:
--threshold THRESHOLD
Prediction threshold (default is 0.35)
--ext EXT Extension to add to caption file in case of dir option (default is .txt)
--overwrite Overwrite the file if it exists.
--overwrite Overwrite the file if it exists
--gpu Use GPU (Compatible CUDA and cuDNN required)
--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}
modelname to use for prediction (default is wd14-convnextv2.v1)
```
@@ -52,3 +54,18 @@ python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2
python run.py --file image.jpg --model mld-caformer.dec-5-97527
python run.py --file image.jpg --model mld-tresnetd.6-30000
```
## Use GPU
Requires CUDA 12.2 and cuDNN8.x.
```
pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
```
```
python run.py --file image.jpg --gpu
```
https://onnxruntime.ai/docs/install/
https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements