Merge branch 'master' into xygrid_infotext_improvements

1 files changed, 37 insertions, 21 deletions

diff --git a/modules/interrogate.py b/modules/interrogate.py
index f62a4745..738d8ff7 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -1,4 +1,3 @@
-import contextlib

 import os

 import sys

 import traceback

@@ -11,25 +10,27 @@ from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode

 

 import modules.shared as shared

-from modules import devices, paths, lowvram

+from modules import devices, paths, lowvram, modelloader

 

 blip_image_eval_size = 384

-blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'

 clip_model_name = 'ViT-L/14'

 

 Category = namedtuple("Category", ["name", "topn", "items"])

 

 re_topn = re.compile(r"\.top(\d+)\.")

 

+

 class InterrogateModels:

     blip_model = None

     clip_model = None

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

@@ -44,7 +45,14 @@ class InterrogateModels:
     def load_blip_model(self):

         import models.blip

 

-        blip_model = models.blip.blip_decoder(pretrained=blip_model_url, image_size=blip_image_eval_size, vit='base', med_config=os.path.join(paths.paths["BLIP"], "configs", "med_config.json"))

+        files = modelloader.load_models(

+            model_path=os.path.join(paths.models_path, "BLIP"),

+            model_url='https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth',

+            ext_filter=[".pth"],

+            download_name='model_base_caption_capfilt_large.pth',

+        )

+

+        blip_model = models.blip.blip_decoder(pretrained=files[0], image_size=blip_image_eval_size, vit='base', med_config=os.path.join(paths.paths["BLIP"], "configs", "med_config.json"))

         blip_model.eval()

 

         return blip_model

@@ -52,26 +60,30 @@ 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", download_root=shared.cmd_opts.clip_models_path)

+        else:

+            model, preprocess = clip.load(clip_model_name, download_root=shared.cmd_opts.clip_models_path)

+

         model.eval()

-        model = model.to(shared.device)

+        model = model.to(devices.device_interrogate)

 

         return model, preprocess

 

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

+            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(shared.device)

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

 

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

 

@@ -98,11 +110,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]

@@ -115,7 +127,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)

@@ -123,8 +135,9 @@ class InterrogateModels:
         return caption[0]

 

     def interrogate(self, pil_image):

-        res = None

-

+        res = ""

+        shared.state.begin()

+        shared.state.job = 'interrogate'

         try:

 

             if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

@@ -139,11 +152,10 @@ class InterrogateModels:
 

             res = caption

 

-            cilp_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"):

-                image_features = self.clip_model.encode_image(cilp_image).type(self.dtype)

+            with torch.no_grad(), devices.autocast():

+                image_features = self.clip_model.encode_image(clip_image).type(self.dtype)

 

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

 

@@ -155,13 +167,17 @@ class InterrogateModels:
                 for name, topn, items in self.categories:

                     matches = self.rank(image_features, items, top_count=topn)

                     for match, score in matches:

-                        res += ", " + match

+                        if shared.opts.interrogate_return_ranks:

+                            res += f", ({match}:{score/100:.3f})"

+                        else:

+                            res += ", " + match

 

         except Exception:

-            print(f"Error interrogating", file=sys.stderr)

+            print("Error interrogating", file=sys.stderr)

             print(traceback.format_exc(), file=sys.stderr)

             res += "<error>"

 

         self.unload()

+        shared.state.end()

 

         return res