diff options
author | yfszzx <yfszzx@gmail.com> | 2022-10-23 16:17:37 +0800 |
---|---|---|
committer | yfszzx <yfszzx@gmail.com> | 2022-10-23 16:17:37 +0800 |
commit | 6a9ea40d7f64139f23d634efd7c2fb2c743a546f (patch) | |
tree | 8d420a5590c09e1d4c005a709166479856e70580 /modules/interrogate.py | |
parent | 67b78f0ea6f196bfdca49932da062631bb40d0b1 (diff) | |
parent | 1ef32c8b8fa3e16a1e7b287eb19d4fc943d1f2a5 (diff) |
Move browser and Inspiration into extension
Diffstat (limited to 'modules/interrogate.py')
-rw-r--r-- | modules/interrogate.py | 12 |
1 files changed, 9 insertions, 3 deletions
diff --git a/modules/interrogate.py b/modules/interrogate.py index 64b91eb4..65b05d34 100644 --- a/modules/interrogate.py +++ b/modules/interrogate.py @@ -28,9 +28,11 @@ class InterrogateModels: clip_preprocess = None
categories = None
dtype = None
+ running_on_cpu = None
def __init__(self, content_dir):
self.categories = []
+ self.running_on_cpu = devices.device_interrogate == torch.device("cpu")
if os.path.exists(content_dir):
for filename in os.listdir(content_dir):
@@ -53,7 +55,11 @@ class InterrogateModels: def load_clip_model(self):
import clip
- model, preprocess = clip.load(clip_model_name)
+ if self.running_on_cpu:
+ model, preprocess = clip.load(clip_model_name, device="cpu")
+ else:
+ model, preprocess = clip.load(clip_model_name)
+
model.eval()
model = model.to(devices.device_interrogate)
@@ -62,14 +68,14 @@ class InterrogateModels: def load(self):
if self.blip_model is None:
self.blip_model = self.load_blip_model()
- if not shared.cmd_opts.no_half:
+ if not shared.cmd_opts.no_half and not self.running_on_cpu:
self.blip_model = self.blip_model.half()
self.blip_model = self.blip_model.to(devices.device_interrogate)
if self.clip_model is None:
self.clip_model, self.clip_preprocess = self.load_clip_model()
- if not shared.cmd_opts.no_half:
+ if not shared.cmd_opts.no_half and not self.running_on_cpu:
self.clip_model = self.clip_model.half()
self.clip_model = self.clip_model.to(devices.device_interrogate)
|