Add 'interrogate' and 'all' choices to --use-cpu

* Add 'interrogate' and 'all' choices to --use-cpu
* Change type for --use-cpu argument to str.lower, so that choices are case insensitive

1 files changed, 7 insertions, 7 deletions

diff --git a/modules/interrogate.py b/modules/interrogate.py
index af858cc0..9263d65a 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -55,7 +55,7 @@ class InterrogateModels:
 

         model, preprocess = clip.load(clip_model_name)

         model.eval()

-        model = model.to(shared.device)

+        model = model.to(devices.device_interrogate)

 

         return model, preprocess

 

@@ -65,14 +65,14 @@ class InterrogateModels:
             if not shared.cmd_opts.no_half:

                 self.blip_model = self.blip_model.half()

 

-        self.blip_model = self.blip_model.to(shared.device)

+        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:

                 self.clip_model = self.clip_model.half()

 

-        self.clip_model = self.clip_model.to(shared.device)

+        self.clip_model = self.clip_model.to(devices.device_interrogate)

 

         self.dtype = next(self.clip_model.parameters()).dtype

 

@@ -99,11 +99,11 @@ class InterrogateModels:
             text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]

 

         top_count = min(top_count, len(text_array))

-        text_tokens = clip.tokenize([text for text in text_array], truncate=True).to(shared.device)

+        text_tokens = clip.tokenize([text for text in text_array], truncate=True).to(devices.device_interrogate)

         text_features = self.clip_model.encode_text(text_tokens).type(self.dtype)

         text_features /= text_features.norm(dim=-1, keepdim=True)

 

-        similarity = torch.zeros((1, len(text_array))).to(shared.device)

+        similarity = torch.zeros((1, len(text_array))).to(devices.device_interrogate)

         for i in range(image_features.shape[0]):

             similarity += (100.0 * image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)

         similarity /= image_features.shape[0]

@@ -116,7 +116,7 @@ class InterrogateModels:
             transforms.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=InterpolationMode.BICUBIC),

             transforms.ToTensor(),

             transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))

-        ])(pil_image).unsqueeze(0).type(self.dtype).to(shared.device)

+        ])(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)

 

         with torch.no_grad():

             caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=shared.opts.interrogate_clip_min_length, max_length=shared.opts.interrogate_clip_max_length)

@@ -140,7 +140,7 @@ class InterrogateModels:
 

             res = caption

 

-            clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(shared.device)

+            clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)

 

             precision_scope = torch.autocast if shared.cmd_opts.precision == "autocast" else contextlib.nullcontext

             with torch.no_grad(), precision_scope("cuda"):