Update app.py

app.py

@@ -108,33 +108,17 @@ def generate_caption_translate(image, target_lang):
 
 
 # small text LM (runs on CPU okay)
-qa_text_model = pipeline("text2text-generation", model="google/flan-t5-large")
-
-def vqa_with_fallback(image, question):
-    # 1) try direct VQA
-    prompt = f"Question: {question} Answer:"
-    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
-    out = model.generate(**inputs, max_new_tokens=40, num_beams=4, early_stopping=True)
-    direct_answer = processor.decode(out[0], skip_special_tokens=True)
-    if direct_answer.lower().startswith(prompt.lower()):
-        direct_answer = direct_answer[len(prompt):].strip()
-
-    # simple heuristics to detect bad/echo answers
-    q_clean = question.strip().lower().rstrip("?.")
-    a_clean = direct_answer.strip().lower().rstrip("?.")
-    bad = (a_clean == "" or a_clean == question.strip().lower() or len(a_clean.split()) <= 2)
-
-    if not bad:
-        return direct_answer
-
-    # 2) fallback: get a caption then use LLM for reasoning
-    cap_inputs = processor(images=image, return_tensors="pt").to(model.device)
-    cap_out = model.generate(**cap_inputs, max_new_tokens=40, num_beams=4)
-    caption = processor.decode(cap_out[0], skip_special_tokens=True)
-
-    # Compose prompt for the text model with grounding
-    text_prompt = f"Image description: {caption}\nQuestion: {question}\nAnswer:"
-    answer = qa_text_model(text_prompt, max_length=80)[0]["generated_text"]
+from transformers import BlipProcessor, BlipForQuestionAnswering
+from PIL import Image
+import torch
+
+vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
+vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to("cuda" if torch.cuda.is_available() else "cpu")
+
+def vqa_proper(image, question):
+    inputs = vqa_processor(image, question, return_tensors="pt").to(vqa_model.device)
+    out = vqa_model.generate(**inputs, max_new_tokens=50, num_beams=5)
+    answer = vqa_processor.decode(out[0], skip_special_tokens=True)
     return answer
 
 
@@ -159,6 +143,6 @@ with gr.Blocks(title="BLIP Vision App") as demo:
             q_in = gr.Textbox(label="Ask a Question about the Image")
         ans_out = gr.Textbox(label="Answer")
         btn2 = gr.Button("Ask")
-        btn2.click(vqa_with_fallback, inputs=[img_vqa, q_in], outputs=ans_out)
+        btn2.click(vqa_proper, inputs=[img_vqa, q_in], outputs=ans_out)
 
 demo.launch()