CLIP interrogator

13 files changed, 204 insertions, 13 deletions

diff --git a/.gitignore b/.gitignore
index 266c8c65..a06b234f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,4 @@ __pycache__
 /embeddings
 /styles.csv
 /webui-user.bat
+/interrogate
diff --git a/README.md b/README.md
index 77b0ff13..24528981 100644
--- a/README.md
+++ b/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)

diff --git a/modules/devices.py b/modules/devices.py
index 25008a04..30d30b99 100644
--- a/modules/devices.py
+++ b/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
diff --git a/modules/interrogate.py b/modules/interrogate.py
new file mode 100644
index 00000000..ed97a58b
--- /dev/null
+++ b/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

diff --git a/modules/paths.py b/modules/paths.py
index 130aecb9..97c17a98 100644
--- a/modules/paths.py
+++ b/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 = {}

diff --git a/modules/shared.py b/modules/shared.py
index 74b0ad89..9eeb64e3 100644
--- a/modules/shared.py
+++ b/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):

diff --git a/modules/ui.py b/modules/ui.py
index 032c20ff..ebc1ae63 100644
--- a/modules/ui.py
+++ b/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,

diff --git a/requirements.txt b/requirements.txt
index ce8f2338..84c9bbd2 100644
--- a/requirements.txt
+++ b/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

diff --git a/requirements_versions.txt b/requirements_versions.txt
index 087040e1..f3a4d2b8 100644
--- a/requirements_versions.txt
+++ b/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

diff --git a/script.js b/script.js
index 9a648021..f41cf90c 100644
--- a/script.js
+++ b/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(){
diff --git a/style.css b/style.css
index 50280000..a57b2420 100644
--- a/style.css
+++ b/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;

 }

diff --git a/webui.bat b/webui.bat
index 1b973448..54734d07 100644
--- a/webui.bat
+++ b/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

diff --git a/webui.py b/webui.py
index 70c68338..ca809f79 100644
--- a/webui.py
+++ b/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()