Update app.py

app.py

@@ -69,15 +69,10 @@
 import gradio as gr
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, pipeline
 from PIL import Image
-import torch
 
-# Load BLIP2 model
+# Load BLIP2 for captioning
 processor = Blip2Processor.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
-)
+blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
 
 # Translation pipelines
 translation_models = {
@@ -86,46 +81,48 @@ translation_models = {
     "Spanish": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
 }
 
-# ---- Caption + Translation ----
-def generate_caption_translate(image, target_lang):
-    inputs = processor(image, return_tensors="pt").to(model.device)
-    out = model.generate(**inputs, max_new_tokens=50)
+# 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)
 
+    # Step 2: Translate caption
     if target_lang in translation_models:
         translated = translation_models[target_lang](english_caption)[0]['translation_text']
     else:
         translated = "Translation not available"
 
-    return english_caption, translated
-
-# ---- 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()
+    # 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, 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()
+