add vqa models to interrogate

2026-09-20 09:38:23 +02:00 · 2024-02-24 14:42:26 -05:00
parent 810fd8b970
commit ffb8f16b7b
4 changed files with 152 additions and 12 deletions
@@ -185,6 +185,10 @@ def create_ui():
                    movement = gr.Label(label="Movement", num_top_classes=5)
                    trending = gr.Label(label="Trending", num_top_classes=5)
                    flavor = gr.Label(label="Flavor", num_top_classes=5)
+                with gr.Row():
+                    clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLIP Model')
+                    ui_common.create_refresh_button(clip_model, get_models, lambda: {"choices": get_models()}, 'refresh_interrogate_models')
+                    mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
                with gr.Row():
                    btn_interrogate_img = gr.Button("Interrogate", variant='primary')
                    btn_analyze_img = gr.Button("Analyze", variant='primary')
@@ -193,6 +197,9 @@ def create_ui():
                    buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "extras", "control"])
                for tabname, button in buttons.items():
                    parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=prompt, source_image_component=image,))
+                btn_interrogate_img.click(interrogate_image, inputs=[image, clip_model, mode], outputs=prompt)
+                btn_analyze_img.click(analyze_image, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
+                btn_unload.click(unload)
            with gr.Tab("Batch"):
                with gr.Row():
                    batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], type='file', interactive=True, height=100)
@@ -204,16 +211,21 @@ def create_ui():
                    batch = gr.Text(label="Prompts", lines=10)
                with gr.Row():
                    write = gr.Checkbox(label='Write prompts to files', value=False)
+                with gr.Row():
+                    clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLIP Model')
+                    ui_common.create_refresh_button(clip_model, get_models, lambda: {"choices": get_models()}, 'refresh_interrogate_models')
                with gr.Row():
                    btn_interrogate_batch = gr.Button("Interrogate", variant='primary')
-        with gr.Column():
-            with gr.Row():
-                # clip_model = gr.Dropdown(get_models(), value='ViT-L-14/openai', label='CLIP Model')
-                clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLIP Model')
-                ui_common.create_refresh_button(clip_model, get_models, lambda: {"choices": get_models()}, 'refresh_interrogate_models')
-            with gr.Row():
-                mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
-        btn_interrogate_img.click(interrogate_image, inputs=[image, clip_model, mode], outputs=prompt)
-        btn_analyze_img.click(analyze_image, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
-        btn_interrogate_batch.click(interrogate_batch, inputs=[batch_files, batch_folder, batch_str, clip_model, mode, write], outputs=[batch])
-        btn_unload.click(unload)
+                btn_interrogate_batch.click(interrogate_batch, inputs=[batch_files, batch_folder, batch_str, clip_model, mode, write], outputs=[batch])
+            with gr.Tab("VQA"):
+                from modules import vqa
+                with gr.Row():
+                    vqa_image = gr.Image(type='pil', label="Image")
+                with gr.Row():
+                    vqa_question = gr.Textbox(label="Question")
+                with gr.Row():
+                    vqa_answer = gr.Textbox(label="Answer", lines=3)
+                with gr.Row():
+                    vqa_model = gr.Dropdown(list(vqa.MODELS), value='None', label='VQA Model')
+                    vqa_submit = gr.Button("Interrogate", variant='primary')
+                vqa_submit.click(vqa.interrogate, inputs=[vqa_question, vqa_image, vqa_model], outputs=[vqa_answer])
@@ -0,0 +1,126 @@
+import torch
+import transformers
+from PIL import Image
+from modules import shared, devices
+
+
+processor = None
+model = None
+loaded: str = None
+MODELS = {
+    "None": None,
+    "GIT TextCaps Base": "microsoft/git-base-textcaps", # 0.7GB
+    "GIT VQA Base": "microsoft/git-base-vqav2", # 0.7GB
+    "GIT VQA Large": "microsoft/git-large-vqav2", # 1.6GB
+    "BLIP Base": "Salesforce/blip-vqa-base", # 1.5GB
+    "BLIP Large": "Salesforce/blip-vqa-capfilt-large", # 1.5GB
+    "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB
+    "Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB
+}
+
+
+def git(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = transformers.GitForCausalLM.from_pretrained(repo)
+        processor = transformers.GitProcessor.from_pretrained(repo)
+        loaded = repo
+    model.to(devices.device, devices.dtype)
+    shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}')
+
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    git_dict = {}
+    git_dict['pixel_values'] = pixel_values.to(devices.device, devices.dtype)
+    if len(question) > 0:
+        input_ids = processor(text=question, add_special_tokens=False).input_ids
+        input_ids = [processor.tokenizer.cls_token_id] + input_ids
+        input_ids = torch.tensor(input_ids).unsqueeze(0)
+        git_dict['input_ids'] = input_ids.to(devices.device)
+    with devices.inference_context():
+        generated_ids = model.generate(**git_dict)
+    response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    model.to(devices.cpu)
+    shared.log.debug(f'VQA: response={response}')
+    return response
+
+
+def blip(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = transformers.BlipForQuestionAnswering.from_pretrained(repo)
+        processor = transformers.BlipProcessor.from_pretrained(repo)
+        loaded = repo
+    model.to(devices.device, devices.dtype)
+    inputs = processor(image, question, return_tensors="pt")
+    inputs = inputs.to(devices.device, devices.dtype)
+    with devices.inference_context():
+        outputs = model.generate(**inputs)
+    response = processor.decode(outputs[0], skip_special_tokens=True)
+
+    model.to(devices.cpu)
+    shared.log.debug(f'VQA: response={response}')
+    return response
+
+
+def vilt(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = transformers.ViltForQuestionAnswering.from_pretrained(repo)
+        processor = transformers.ViltProcessor.from_pretrained(repo)
+        loaded = repo
+    model.to(devices.device)
+    shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}')
+
+    inputs = processor(image, question, return_tensors="pt")
+    inputs = inputs.to(devices.device)
+    with devices.inference_context():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    idx = logits.argmax(-1).item()
+    response = model.config.id2label[idx]
+
+    model.to(devices.cpu)
+    shared.log.debug(f'VQA: response={response}')
+    return response
+
+
+def pix(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = transformers.Pix2StructForConditionalGeneration.from_pretrained(repo)
+        processor = transformers.Pix2StructProcessor.from_pretrained(repo)
+        loaded = repo
+    model.to(devices.device)
+    shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}')
+
+    if len(question) > 0:
+        inputs = processor(images=image, text=question, return_tensors="pt").to(devices.device)
+    else:
+        inputs = processor(images=image, return_tensors="pt").to(devices.device)
+    with devices.inference_context():
+        outputs = model.generate(**inputs)
+    response = processor.decode(outputs[0], skip_special_tokens=True)
+
+    model.to(devices.cpu)
+    shared.log.debug(f'VQA: response={response}')
+    return response
+
+
+def interrogate(vqa_question, vqa_image, vqa_model):
+    vqa_model = MODELS.get(vqa_model, None)
+    shared.log.debug(f'VQA: model="{vqa_model}" question={vqa_question} image={vqa_image}')
+    if vqa_image is None:
+        return 'no image provided'
+    if vqa_model is None:
+        return 'no model selected'
+    if 'git' in vqa_model.lower():
+        return git(vqa_question, vqa_image, vqa_model)
+    if 'vilt' in vqa_model.lower():
+        return vilt(vqa_question, vqa_image, vqa_model)
+    if 'blip' in vqa_model.lower():
+        return blip(vqa_question, vqa_image, vqa_model)
+    if 'pix' in vqa_model.lower():
+        return pix(vqa_question, vqa_image, vqa_model)
+    else:
+        return 'unknown model'