Support gpu.

README.md

@@ -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

requirements.txt

@@ -1,7 +1,6 @@
 pandas==2.1.1
 huggingface-hub==0.17.3
-torch==2.0.1
 Pillow==10.0.1
 deepdanbooru==1.0.2
-onnxruntime==1.16.0
+onnxruntime==1.17.3
 opencv-python==4.8.0.76

run.py

@@ -1,6 +1,8 @@
 
 
 import sys
+
+from sympy import use
 from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
 from PIL import Image
 from pathlib import Path
@@ -27,6 +29,10 @@ parser.add_argument(
     '--overwrite',
     action='store_true',
     help='Overwrite the file if it exists')
+parser.add_argument(
+    '--gpu',
+    action='store_true',
+    help='Use GPU (Compatible CUDA and cuDNN required)')
 parser.add_argument(
     '--model',
     default='wd14-convnextv2.v1',
@@ -37,6 +43,9 @@ args = parser.parse_args()
 # get interrogator configs
 interrogator = interrogators[args.model]
 
+if args.gpu:
+    interrogator.use_gpu()
+
 def image_interrogate(image_path: Path):
     """
     Predictions from a image path

tagger/interrogator.py

@@ -17,7 +17,6 @@ tag_escape_pattern = re.compile(r'([\\()])')
 
 import tagger.dbimutils as dbimutils
 
-use_cpu = True
 
 class Interrogator:
     @staticmethod
@@ -73,6 +72,7 @@ class Interrogator:
 
     def __init__(self, name: str) -> None:
         self.name = name
+        self.providers = ['CPUExecutionProvider']
 
     def load(self):
         raise NotImplementedError()
@@ -90,6 +90,9 @@ class Interrogator:
 
         return unloaded
 
+    def use_gpu(self) -> None:
+        self.providers = ['CUDAExecutionProvider']
+
     def interrogate(
         self,
         image: Image
@@ -111,6 +114,7 @@ class WaifuDiffusionInterrogator(Interrogator):
         self.model_path = model_path
         self.tags_path = tags_path
         self.kwargs = kwargs
+        self.use_cpu = True
 
     def download(self) -> Tuple[os.PathLike, os.PathLike]:
         print(f"Loading {self.name} model file from {self.kwargs['repo_id']}")
@@ -125,14 +129,7 @@ class WaifuDiffusionInterrogator(Interrogator):
         model_path, tags_path = self.download()
 
         from onnxruntime import InferenceSession
-
-        # https://onnxruntime.ai/docs/execution-providers/
-        # https://github.com/toriato/stable-diffusion-webui-wd14-tagger/commit/e4ec460122cf674bbf984df30cdb10b4370c1224#r92654958
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-        if use_cpu:
-            providers.pop(0)
-
-        self.model = InferenceSession(str(model_path), providers=providers)
+        self.model = InferenceSession(str(model_path), providers=self.providers)
 
         print(f'Loaded {self.name} model from {model_path}')
 
@@ -202,6 +199,7 @@ class MLDanbooruInterrogator(Interrogator):
         self.repo_id = repo_id
         self.tags = None
         self.model = None
+        self.use_cpu = False
 
     def download(self) -> Tuple[str, str]:
         print(f"Loading {self.name} model file from {self.repo_id}")
@@ -220,11 +218,8 @@ class MLDanbooruInterrogator(Interrogator):
         model_path, tags_path = self.download()
 
         from onnxruntime import InferenceSession
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-        if use_cpu:
-            providers.pop(0)
         self.model = InferenceSession(model_path,
-                                          providers=providers)
+                                        providers=self.providers)
         print(f'Loaded {self.name} model from {model_path}')
 
         with open(tags_path, 'r', encoding='utf-8') as filen: