Merge pull request #4 from corkborg/gpu

Support GPU.

README.md

@@ -13,7 +13,7 @@ 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] [--cpu] [--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 +22,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 caption file if it exists
+  --cpu                 Use CPU only
   --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 +53,14 @@ 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
 ```
+
+## Using 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/
+```
+
+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
+pandas==2.2.2
+huggingface-hub==0.22.2
 Pillow==10.0.1
 deepdanbooru==1.0.2
-onnxruntime==1.16.0
-opencv-python==4.8.0.76
+onnxruntime==1.17.3
+opencv-python==4.9.0.80

run.py

@@ -1,6 +1,3 @@
-
-
-import sys
 from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
 from PIL import Image
 from pathlib import Path
@@ -26,7 +23,11 @@ parser.add_argument(
 parser.add_argument(
     '--overwrite',
     action='store_true',
-    help='Overwrite the file if it exists')
+    help='Overwrite caption file if it exists')
+parser.add_argument(
+    '--cpu',
+    action='store_true',
+    help='Use CPU only')
 parser.add_argument(
     '--model',
     default='wd14-convnextv2.v1',
@@ -37,6 +38,9 @@ args = parser.parse_args()
 # get interrogator configs
 interrogator = interrogators[args.model]
 
+if args.cpu:
+    interrogator.use_cpu()
+
 def image_interrogate(image_path: Path):
     """
     Predictions from a image path
@@ -63,7 +67,7 @@ if args.dir:
         tags = image_interrogate(image_path)
         tags_str = ", ".join(tags.keys())
 
-        with open(caption_path, "w") as fp:
+        with open(caption_path, 'w') as fp:
             fp.write(tags_str)
 
 if args.file:

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 = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 
     def load(self):
         raise NotImplementedError()
@@ -90,6 +90,9 @@ class Interrogator:
 
         return unloaded
 
+    def use_cpu(self) -> None:
+        self.providers = ['CPUExecutionProvider']
+
     def interrogate(
         self,
         image: Image
@@ -125,14 +128,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}')
 
@@ -220,11 +216,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: