Support gpu.

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: