refactor: app

app.py

@@ -1,25 +1,28 @@
+from PIL import Image
 import gradio as gr
-import spaces
-import torch
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+import requests
+from transformers import AutoProcessor, BlipForQuestionAnswering
 
-model_id = "llava-hf/llava-1.5-7b-hf"
+model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
+processor = AutoProcessor.from_pretrained("Salesforce/blip-vqa-base")
 
-prompt_format = "USER: <image>\n{}\nASSISTANT:"
+def generate_answer(text, image):
 
-model = LlavaForConditionalGeneration.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float16, 
-    low_cpu_mem_usage=True,
-).cuda()
+    text = text
+    inputs = processor(images=image, text=text, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    
+    return processor.decode(outputs[0], skip_special_tokens=True)
 
-processor = AutoProcessor.from_pretrained(model_id)
+text_input = gr.Textbox(lines=5, label="Enter text")
+image_input = gr.Image(type="pil", label="Upload Image")
 
-@spaces.GPU
-def inference(text, image):
-    prompt = prompt_format.format(text)
-    inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
-    output = model.generate(**inputs, max_new_tokens=1024)
-    return processor.decode(output[0], skip_special_tokens=True).split("ASSISTANT:")[-1]
+iface = gr.Interface(
+    fn=generate_answer, 
+    inputs=[text_input, image_input], 
+    outputs="text",
+    title="DD360-Bot-Multimodal",
+    description="Enter text and upload an image"
+)
 
-gr.Interface(fn=inference, inputs=[gr.Text(), gr.Image()], outputs=gr.Text()).launch()
+iface.launch()