Add cpu only option instead of gpu option.

README.md

@@ -24,7 +24,7 @@ options:
                         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
-  --gpu                 Use GPU (Compatible CUDA and cuDNN required)
+  --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)
 ```
@@ -63,9 +63,5 @@ 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

run.py

@@ -1,8 +1,3 @@
-
-
-import sys
-
-from sympy import use
 from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
 from PIL import Image
 from pathlib import Path
@@ -30,9 +25,9 @@ parser.add_argument(
     action='store_true',
     help='Overwrite the file if it exists')
 parser.add_argument(
-    '--gpu',
+    '--cpu',
     action='store_true',
-    help='Use GPU (Compatible CUDA and cuDNN required)')
+    help='Use CPU only')
 parser.add_argument(
     '--model',
     default='wd14-convnextv2.v1',
@@ -43,8 +38,8 @@ args = parser.parse_args()
 # get interrogator configs
 interrogator = interrogators[args.model]
 
-if args.gpu:
-    interrogator.use_gpu()
+if args.cpu:
+    interrogator.use_cpu()
 
 def image_interrogate(image_path: Path):
     """
@@ -72,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

@@ -72,7 +72,7 @@ class Interrogator:
 
     def __init__(self, name: str) -> None:
         self.name = name
-        self.providers = ['CPUExecutionProvider']
+        self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 
     def load(self):
         raise NotImplementedError()
@@ -90,8 +90,8 @@ class Interrogator:
 
         return unloaded
 
-    def use_gpu(self) -> None:
-        self.providers = ['CUDAExecutionProvider']
+    def use_cpu(self) -> None:
+        self.providers = ['CPUExecutionProvider']
 
     def interrogate(
         self,