Update app.py

app.py

@@ -69,10 +69,15 @@
 import gradio as gr
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, pipeline
 from PIL import Image
+import torch
 
-# Load BLIP model
+# Load BLIP2 model
 processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
+model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-opt-2.7b",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None
+)
 
 # Translation pipelines
 translation_models = {
@@ -81,13 +86,12 @@ translation_models = {
     "Spanish": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
 }
 
+# ---- Caption + Translation ----
 def generate_caption_translate(image, target_lang):
-    # Step 1: Generate English caption
-    inputs = processor(image, return_tensors="pt")
+    inputs = processor(image, return_tensors="pt").to(model.device)
     out = model.generate(**inputs, max_new_tokens=50)
     english_caption = processor.decode(out[0], skip_special_tokens=True)
 
-    # Step 2: Translate
     if target_lang in translation_models:
         translated = translation_models[target_lang](english_caption)[0]['translation_text']
     else:
@@ -95,18 +99,34 @@ def generate_caption_translate(image, target_lang):
 
     return english_caption, translated
 
-# Gradio Interface
-interface = gr.Interface(
-    fn=generate_caption_translate,
-    inputs=[gr.Image(type="pil"), gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To")],
-    outputs=[
-        gr.Textbox(label="English Caption"),
-        gr.Textbox(label="Translated Caption")
-    ],
-    title="BLIP Captioning + Translation"
-)
+# ---- Visual Question Answering ----
+def answer_question(image, question):
+    inputs = processor(image, text=question, return_tensors="pt").to(model.device)
+    out = model.generate(**inputs, max_new_tokens=50)
+    answer = processor.decode(out[0], skip_special_tokens=True)
+    return answer
+
+# ---- Gradio Interface ----
+with gr.Blocks() as demo:
+    gr.Markdown("## 🖼️ BLIP2: Image Captioning + Translation + VQA")
+
+    with gr.Tab("Caption + Translation"):
+        img1 = gr.Image(type="pil")
+        lang = gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To")
+        eng_cap = gr.Textbox(label="English Caption")
+        trans_cap = gr.Textbox(label="Translated Caption")
+        btn1 = gr.Button("Generate Caption + Translate")
+        btn1.click(generate_caption_translate, inputs=[img1, lang], outputs=[eng_cap, trans_cap])
+
+    with gr.Tab("Visual Question Answering"):
+        img2 = gr.Image(type="pil")
+        question = gr.Textbox(label="Ask a Question about the Image")
+        answer = gr.Textbox(label="Answer")
+        btn2 = gr.Button("Get Answer")
+        btn2.click(answer_question, inputs=[img2, question], outputs=answer)
+
+demo.launch()
 
-interface.launch()