Update app.py

app.py

@@ -69,59 +69,80 @@
 import gradio as gr
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, pipeline
 from PIL import Image
+import torch
+
+# ----------------------
+# Load BLIP2 for Captioning
+# ----------------------
+caption_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+caption_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
+
+# ----------------------
+# Load BLIP2 for VQA
+# ----------------------
+vqa_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
+vqa_model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-flan-t5-xl", torch_dtype=torch.float16, device_map="auto"
+)
 
-# Load BLIP2 for captioning
-processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
-
+# ----------------------
 # Translation pipelines
+# ----------------------
 translation_models = {
     "Hindi": pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi"),
     "French": pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr"),
     "Spanish": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
 }
 
-# Language model for reasoning/Q&A
-qa_model = pipeline("text2text-generation", model="google/flan-t5-large")
-
-def caption_translate_vqa(image, target_lang, question):
-    # Step 1: Generate English caption
-    inputs = processor(image, return_tensors="pt")
-    out = blip_model.generate(**inputs, max_new_tokens=50)
-    english_caption = processor.decode(out[0], skip_special_tokens=True)
+# ----------------------
+# Caption + Translate Function
+# ----------------------
+def generate_caption_translate(image, target_lang):
+    inputs = caption_processor(image, return_tensors="pt")
+    out = caption_model.generate(**inputs, max_new_tokens=50)
+    english_caption = caption_processor.decode(out[0], skip_special_tokens=True)
 
-    # Step 2: Translate caption
+    # Translate if chosen
     if target_lang in translation_models:
         translated = translation_models[target_lang](english_caption)[0]['translation_text']
     else:
         translated = "Translation not available"
 
-    # Step 3: Image Q&A using caption + question
-    if question and len(question.strip()) > 0:
-        prompt = f"Image description: {english_caption}\nQuestion: {question}\nAnswer:"
-        answer = qa_model(prompt, max_length=100)[0]['generated_text']
-    else:
-        answer = "No question asked."
+    return english_caption, translated
 
-    return english_caption, translated, answer
+# ----------------------
+# VQA Function
+# ----------------------
+def vqa(image, question):
+    inputs = vqa_processor(image, question, return_tensors="pt").to(vqa_model.device)
+    out = vqa_model.generate(**inputs, max_new_tokens=100)
+    answer = vqa_processor.decode(out[0], skip_special_tokens=True)
+    return answer
 
+# ----------------------
 # Gradio UI
-interface = gr.Interface(
-    fn=caption_translate_vqa,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To"),
-        gr.Textbox(label="Ask a Question about the Image")
-    ],
-    outputs=[
-        gr.Textbox(label="English Caption"),
-        gr.Textbox(label="Translated Caption"),
-        gr.Textbox(label="VQA Answer")
-    ],
-    title="BLIP2 + Translation + Visual Q&A"
-)
-
-interface.launch()
+# ----------------------
+with gr.Blocks(title="BLIP2 Vision App") as demo:
+    gr.Markdown("## 🖼️ BLIP2: Image Captioning + Translation + Question Answering")
+
+    with gr.Tab("Caption + Translate"):
+        with gr.Row():
+            img_in = gr.Image(type="pil", label="Upload Image")
+            lang_in = gr.Dropdown(["Hindi", "French", "Spanish"], label="Translate To")
+        eng_out = gr.Textbox(label="English Caption")
+        trans_out = gr.Textbox(label="Translated Caption")
+        btn1 = gr.Button("Generate Caption & Translate")
+        btn1.click(generate_caption_translate, inputs=[img_in, lang_in], outputs=[eng_out, trans_out])
+
+    with gr.Tab("Visual Question Answering (VQA)"):
+        with gr.Row():
+            img_vqa = gr.Image(type="pil", label="Upload Image")
+            q_in = gr.Textbox(label="Ask a Question about the Image")
+        ans_out = gr.Textbox(label="Answer")
+        btn2 = gr.Button("Ask")
+        btn2.click(vqa, inputs=[img_vqa, q_in], outputs=ans_out)
+
+demo.launch()