Update app.py

app.py

@@ -65,12 +65,13 @@ app = gr.mount_gradio_app(app, gui, path="/")
 @app.get("/")
 def home():
     return RedirectResponse(url="/") """
-from fastapi import FastAPI, UploadFile, Form
-from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
-import os
-import shutil
+from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
 from PIL import Image
-from transformers import ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    ViltProcessor, ViltForQuestionAnswering,
+    T5Tokenizer, T5ForConditionalGeneration
+)
 from gtts import gTTS
 import torch
 import tempfile
@@ -78,28 +79,20 @@ import gradio as gr
 
 app = FastAPI()
 
-# Load VQA Model
+# VQA Model
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 
-# Load Falcon-7B-Instruct model to rewrite answers
-gpt_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
-gpt_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct")
+# Text Rewriter (FLAN-T5-base)
+rewrite_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
+rewrite_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
 
 def rewrite_answer(question, short_answer):
-    prompt = f"Rewrite the short answer into a natural sentence.\nQuestion: {question}\nAnswer: {short_answer}\nFull Sentence:"
-    inputs = gpt_tokenizer(prompt, return_tensors="pt")
+    prompt = f"Answer the question '{question}' with a complete sentence using this answer: '{short_answer}'"
+    inputs = rewrite_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
-        outputs = gpt_model.generate(
-            **inputs,
-            max_new_tokens=50,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.8,
-            pad_token_id=gpt_tokenizer.eos_token_id
-        )
-    rewritten = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return rewritten.split("Full Sentence:")[-1].strip()
+        outputs = rewrite_model.generate(**inputs, max_new_tokens=50)
+    return rewrite_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def answer_question_from_image(image, question):
     if image is None or not question.strip():
@@ -111,23 +104,21 @@ def answer_question_from_image(image, question):
     predicted_id = outputs.logits.argmax(-1).item()
     short_answer = vqa_model.config.id2label[predicted_id]
 
-    # Rewrite short answer to full sentence
+    # Rewrite to full sentence
     full_answer = rewrite_answer(question, short_answer)
 
     try:
         tts = gTTS(text=full_answer)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
             tts.save(tmp.name)
-            audio_path = tmp.name
+            return full_answer, tmp.name
     except Exception as e:
-        return f"Answer: {full_answer}\n\n⚠️ Audio generation error: {e}", None
-
-    return full_answer, audio_path
+        return f"{full_answer}\n\n⚠️ Audio generation error: {e}", None
 
 def process_image_question(image: Image.Image, question: str):
-    answer, audio_path = answer_question_from_image(image, question)
-    return answer, audio_path
+    return answer_question_from_image(image, question)
 
+# Gradio UI
 gui = gr.Interface(
     fn=process_image_question,
     inputs=[
@@ -146,4 +137,4 @@ app = gr.mount_gradio_app(app, gui, path="/")
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")