using 3 new models of VQA

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import os
 
-from inference import *
+from inference import Inference
 from utils import *
 
+inference = Inference()
 
 with gr.Blocks() as block:
     options = gr.Dropdown(choices=["Model 1", "Model 2", "Model 3"], label="Models", info="Select the model to use..")
@@ -15,9 +16,16 @@ with gr.Blocks() as block:
     dogs = os.path.join(os.path.dirname(__file__), "617.jpg")
     image = gr.Image(type="pil", value=dogs)
 
-    btn.click(inference, inputs=[image, txt], outputs=[txt_3])
+    btn.click(inference.inference_vilt, inputs=[image, txt], outputs=[txt_3])
 
     btn = gr.Button(value="Submit")
 
+iface = gr.Interface(
+    fn=lambda: "Selected Model: " + options.value,
+    inputs=block,
+    outputs="text",
+    live=False
+)
+
 if __name__ == "__main__":
-    block.launch()
+    iface.launch()

inference.py

@@ -1,11 +1,32 @@
-from transformers import ViltProcessor, ViltForQuestionAnswering
+from transformers import ViltProcessor, ViltForQuestionAnswering, Pix2StructProcessor, Pix2StructForConditionalGeneration, Blip2Processor, Blip2ForConditionalGeneration
+import torch
 
-def inference(image, text):
-    processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-    model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-    encoding = processor(image, text, return_tensors="pt")
+class Inference:
+    def __init__(self):
+        self.vilt_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+        self.vilt_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 
-    outputs = model(**encoding)
-    logits = outputs.logits
-    idx = logits.argmax(-1).item()
-    return f"{model.config.id2label[idx]}"
+        self.deplot_processor = Pix2StructProcessor.from_pretrained('google/deplot')
+        self.deplot_model = Pix2StructForConditionalGeneration.from_pretrained('google/deplot')
+
+        self.blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+        self.blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map={"": 0}, torch_dtype=torch.float16)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def inference_vilt(self, image, text):
+        encoding = self.vilt_processor(image, text, return_tensors="pt")
+        outputs = self.vilt_model(**encoding)
+        logits = outputs.logits
+        idx = logits.argmax(-1).item()
+        return f"{self.vilt_model.config.id2label[idx]}"
+
+    def inference_deplot(self, image, text):
+        inputs = self.deplot_processor(images=image, text=text, return_tensors="pt")
+        predictions = self.deplot_model.generate(**inputs, max_new_tokens=512)
+        return f"{self.deplot_processor.decode(predictions[0], skip_special_tokens=True)}"
+
+    def inference_vilt(self, image, text):
+        inputs = self.blip_processor(images=image, return_tensors="pt").to(self.device, torch.float16)
+        generated_ids = self.blip_model.generate(**inputs)
+        generated_text = self.blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        return f"{generated_text}"