CLIP interrogator

.gitignore

@@ -13,3 +13,4 @@ __pycache__
 /embeddings
 /styles.csv
 /webui-user.bat
+/interrogate

README.md

@@ -40,6 +40,7 @@ A browser interface based on Gradio library for Stable Diffusion.
 - Styles
 - Variations
 - Seed resizing
+- CLIP interrogator
 
 ## Installing and running
 
@@ -289,5 +290,6 @@ After that follow the instructions in the `Manual instructions` section starting
 - Ideas for optimizations - https://github.com/basujindal/stable-diffusion
 - Cross Attention layer optimization - https://github.com/Doggettx/stable-diffusion
 - Idea for SD upscale - https://github.com/jquesnelle/txt2imghd
+- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator
 - Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user.
 - (You)

modules/devices.py

@@ -1,12 +1,16 @@
 import torch
 
-
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
 has_mps = getattr(torch, 'has_mps', False)
 
+cpu = torch.device("cpu")
+
+
 def get_optimal_device():
-  if torch.cuda.is_available():
-      return torch.device("cuda")
-  if has_mps:
-      return torch.device("mps")
-  return torch.device("cpu")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    if has_mps:
+        return torch.device("mps")
+
+    return cpu

modules/interrogate.py

@@ -0,0 +1,142 @@
+import os
+import sys
+import traceback
+from collections import namedtuple
+import re
+
+import torch
+
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+
+import modules.shared as shared
+from modules import devices, paths
+
+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
+
+    def __init__(self, content_dir):
+        self.categories = []
+
+        if os.path.exists(content_dir):
+            for filename in os.listdir(content_dir):
+                m = re_topn.search(filename)
+                topn = 1 if m is None else int(m.group(1))
+
+                with open(os.path.join(content_dir, filename), "r", encoding="utf8") as file:
+                    lines = [x.strip() for x in file.readlines()]
+
+                self.categories.append(Category(name=filename, topn=topn, items=lines))
+
+    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"))
+        blip_model.eval()
+
+        return blip_model
+
+    def load_clip_model(self):
+        import clip
+
+        model, preprocess = clip.load(clip_model_name)
+        model.eval()
+        model = model.to(shared.device)
+
+        return model, preprocess
+
+    def load(self):
+        if self.blip_model is None:
+            self.blip_model = self.load_blip_model()
+
+        self.blip_model = self.blip_model.to(shared.device)
+
+        if self.clip_model is None:
+            self.clip_model, self.clip_preprocess = self.load_clip_model()
+
+        self.clip_model = self.clip_model.to(shared.device)
+
+    def unload(self):
+        if not shared.opts.interrogate_keep_models_in_memory:
+            if self.clip_model is not None:
+                self.clip_model = self.clip_model.to(devices.cpu)
+
+            if self.blip_model is not None:
+                self.blip_model = self.blip_model.to(devices.cpu)
+
+
+    def rank(self, image_features, text_array, top_count=1):
+        import clip
+
+        top_count = min(top_count, len(text_array))
+        text_tokens = clip.tokenize([text for text in text_array]).cuda()
+        with torch.no_grad():
+            text_features = self.clip_model.encode_text(text_tokens).float()
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+
+        similarity = torch.zeros((1, len(text_array))).to(shared.device)
+        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]
+
+        top_probs, top_labels = similarity.cpu().topk(top_count, dim=-1)
+        return [(text_array[top_labels[0][i].numpy()], (top_probs[0][i].numpy()*100)) for i in range(top_count)]
+
+
+    def generate_caption(self, pil_image):
+        gpu_image = transforms.Compose([
+            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).to(shared.device)
+
+        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)
+
+        return caption[0]
+
+    def interrogate(self, pil_image):
+        res = None
+
+        try:
+            self.load()
+
+            caption = self.generate_caption(pil_image)
+            res = caption
+
+            images = self.clip_preprocess(pil_image).unsqueeze(0).to(shared.device)
+
+            with torch.no_grad():
+                image_features = self.clip_model.encode_image(images).float()
+
+            image_features /= image_features.norm(dim=-1, keepdim=True)
+
+            if shared.opts.interrogate_use_builtin_artists:
+                artist = self.rank(image_features, ["by " + artist.name for artist in shared.artist_db.artists])[0]
+
+                res += ", " + artist[0]
+
+            for name, topn, items in self.categories:
+                matches = self.rank(image_features, items, top_count=topn)
+                for match, score in matches:
+                    res += ", " + match
+
+        except Exception:
+            print(f"Error interrogating", file=sys.stderr)
+            print(traceback.format_exc(), file=sys.stderr)
+
+        self.unload()
+
+        return res

modules/paths.py

@@ -18,6 +18,7 @@ path_dirs = [
     (sd_path, 'ldm', 'Stable Diffusion'),
     (os.path.join(sd_path, '../taming-transformers'), 'taming', 'Taming Transformers'),
     (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer'),
+    (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP'),
 ]
 
 paths = {}

modules/shared.py

@@ -11,6 +11,7 @@ import modules.artists
 from modules.paths import script_path, sd_path
 from modules.devices import get_optimal_device
 import modules.styles
+import modules.interrogate
 
 config_filename = "config.json"
 
@@ -77,6 +78,8 @@ artist_db = modules.artists.ArtistsDatabase(os.path.join(script_path, 'artists.c
 styles_filename = os.path.join(script_path, 'styles.csv')
 prompt_styles = modules.styles.load_styles(styles_filename)
 
+interrogator = modules.interrogate.InterrogateModels("interrogate")
+
 face_restorers = []
 
 class Options:
@@ -123,6 +126,11 @@ class Options:
         "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job. Broken in PyCharm console."),
         "face_restoration_model": OptionInfo(None, "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}),
         "code_former_weight": OptionInfo(0.5, "CodeFormer weight parameter; 0 = maximum effect; 1 = minimum effect", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
+        "interrogate_keep_models_in_memory": OptionInfo(True, "Interrogate: keep models in VRAM"),
+        "interrogate_use_builtin_artists": OptionInfo(True, "Interrogate: use artists from artists.csv"),
+        "interrogate_clip_num_beams": OptionInfo(1, "Interrogate: num_beams for BLIP", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}),
+        "interrogate_clip_min_length": OptionInfo(24, "Interrogate: minimum descripton length (excluding artists, etc..)", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}),
+        "interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum descripton length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),
     }
 
     def __init__(self):

modules/ui.py

@@ -242,9 +242,14 @@ def add_style(style_name, text):
     return [update, update]
 
 
+def interrogate(image):
+    prompt = shared.interrogator.interrogate(image)
+
+    return gr_show(True) if prompt is None else prompt
+
 def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     with gr.Blocks(analytics_enabled=False) as txt2img_interface:
-        with gr.Row():
+        with gr.Row(elem_id="toprow"):
             txt2img_prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1)
             negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
             txt2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
@@ -365,10 +370,11 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             )
 
     with gr.Blocks(analytics_enabled=False) as img2img_interface:
-        with gr.Row():
+        with gr.Row(elem_id="toprow"):
             img2img_prompt = gr.Textbox(label="Prompt", elem_id="img2img_prompt", show_label=False, placeholder="Prompt", lines=1)
             negative_prompt = gr.Textbox(label="Negative prompt", elem_id="img2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
             img2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
+            img2img_interrogate = gr.Button('Interrogate', elem_id="img2img_interrogate", variant='primary')
             submit = gr.Button('Generate', elem_id="img2img_generate", variant='primary')
             check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
 
@@ -461,6 +467,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     inpaint_full_res: gr_show(is_inpaint),
                     inpainting_mask_invert: gr_show(is_inpaint),
                     denoising_strength_change_factor: gr_show(is_loopback),
+                    img2img_interrogate: gr_show(not is_inpaint),
                 }
 
             switch_mode.change(
@@ -480,6 +487,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     inpaint_full_res,
                     inpainting_mask_invert,
                     denoising_strength_change_factor,
+                    img2img_interrogate,
                 ]
             )
 
@@ -540,6 +548,12 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             img2img_prompt.submit(**img2img_args)
             submit.click(**img2img_args)
 
+            img2img_interrogate.click(
+                fn=interrogate,
+                inputs=[init_img],
+                outputs=[img2img_prompt],
+            )
+
             check_progress.click(
                 fn=check_progress_call,
                 show_progress=False,

requirements.txt

@@ -15,3 +15,5 @@ fonts
 font-roboto
 git+https://github.com/crowsonkb/k-diffusion.git
 git+https://github.com/TencentARC/GFPGAN.git
+timm==0.4.12
+fairscale==0.4.4

requirements_versions.txt

@@ -11,3 +11,5 @@ pytorch_lightning==1.7.2
 scikit-image==0.19.2
 fonts
 font-roboto
+timm==0.4.12
+fairscale==0.4.4

script.js

@@ -51,6 +51,8 @@ titles = {
     "Variation strength": "How strong of a variation to produce. At 0, there will be no effect. At 1, you will get the complete picture with variation seed (except for ancestral samplers, where you will just get something).",
     "Resize seed from height": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution",
     "Resize seed from width": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution",
+
+    "Interrogate": "Reconstruct frompt from existing image and put it into the prompt field.",
 }
 
 function gradioApp(){

style.css

@@ -5,6 +5,10 @@
     max-width: 13em;
 }
 
+#img2img_interrogate{
+    max-width: 10em;
+}
+
 #subseed_show{
     min-width: 6em;
     max-width: 6em;
@@ -26,7 +30,7 @@
     padding-right: 0;
 }
 
-#component-1 div{
+#toprow div{
     border: none;
     gap: 0;
 }

webui.bat

@@ -85,7 +85,7 @@ if %ERRORLEVEL% == 0 goto :install_reqs
 goto :show_stdout_stderr
 
 :install_reqs
-%PYTHON% -c "import omegaconf; import fonts" >tmp/stdout.txt 2>tmp/stderr.txt
+%PYTHON% -c "import omegaconf; import fonts; import timm" >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :make_dirs
 echo Installing requirements...
 %PYTHON% -m pip install -r %REQS_FILE% --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
@@ -117,12 +117,19 @@ goto :show_stdout_stderr
 
 :install_codeformer_reqs
 %PYTHON% -c "import lpips" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_model
+if %ERRORLEVEL% == 0 goto :clone_blip
 echo Installing requirements for CodeFormer...
 %PYTHON% -m pip install -r repositories\CodeFormer\requirements.txt --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_model
+if %ERRORLEVEL% == 0 goto :clone_blip
 goto :show_stdout_stderr
 
+:clone_blip
+if exist repositories\BLIP goto :check_model
+echo Cloning BLIP repository...
+%GIT% clone https://github.com/salesforce/BLIP.git repositories\BLIP >tmp/stdout.txt 2>tmp/stderr.txt
+if %ERRORLEVEL% NEQ 0 goto :show_stdout_stderr
+%GIT% -C repositories/BLIP checkout 48211a1594f1321b00f14c9f7a5b4813144b2fb9 >tmp/stdout.txt 2>tmp/stderr.txt
+if %ERRORLEVEL% NEQ 0 goto :show_stdout_stderr
 
 :check_model
 dir model.ckpt >tmp/stdout.txt 2>tmp/stderr.txt

webui.py

@@ -20,7 +20,7 @@ import modules.gfpgan_model
 import modules.face_restoration
 import modules.realesrgan_model as realesrgan
 import modules.esrgan_model as esrgan
-import  modules.extras
+import modules.extras
 import modules.lowvram
 import modules.txt2img
 import modules.img2img
@@ -33,6 +33,7 @@ shared.face_restorers.append(modules.face_restoration.FaceRestoration())
 esrgan.load_models(cmd_opts.esrgan_models_path)
 realesrgan.setup_realesrgan()
 
+
 def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
     pl_sd = torch.load(ckpt, map_location="cpu")
@@ -116,5 +117,6 @@ def webui():
 
     demo.launch(share=cmd_opts.share, server_name="0.0.0.0" if cmd_opts.listen else None, server_port=cmd_opts.port)
 
+
 if __name__ == "__main__":
     webui()