Update app.py

app.py

@@ -1,22 +1,22 @@
 from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import RedirectResponse
 import fitz  # PyMuPDF for PDF parsing
 from tika import parser  # Apache Tika for document parsing
 import openpyxl
 from pptx import Presentation
-import torch
+from PIL import Image
 from transformers import pipeline
 import gradio as gr
-from PIL import Image
 import numpy as np
 
-# Initialize FastAPI (not needed for HF Spaces but kept for flexibility)
+# Initialize FastAPI
 app = FastAPI()
 
 print(f"🔄 Loading models")
 
 # Load Hugging Face Models
 doc_qa_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)
-vqa_pipeline = pipeline("vqa", model="Salesforce/blip-vqa-base")  # VQA model for images
+image_captioning_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
 print("✅ Models loaded")
 
@@ -67,7 +67,6 @@ def extract_text_from_excel(excel_bytes):
     except Exception as e:
         return f"❌ Error reading Excel: {str(e)}"
 
-# Function to process documents and answer questions
 def answer_question_from_document(file: UploadFile, question: str):
     print("📂 Processing document for QA...")
     validation_error = validate_file_type(file)
@@ -95,17 +94,19 @@ def answer_question_from_document(file: UploadFile, question: str):
     
     return response[0]["generated_text"]
 
-# Function to process images and answer questions (NO OCR)
 def answer_question_from_image(image, question):
     try:
         print("🖼️ Processing image for QA...")
         if isinstance(image, np.ndarray):  # If it's a NumPy array from Gradio
             image = Image.fromarray(image)  # Convert to PIL Image
         
-        print("🤖 Answering question based on image content...")
-        response = vqa_pipeline(image=image, question=question)
+        print("🖼️ Generating caption for image...")
+        caption = image_captioning_pipeline(image)[0]['generated_text']
+
+        print("🤖 Answering question based on caption...")
+        response = doc_qa_pipeline(f"Question: {question}\nContext: {caption}")
 
-        return response[0]["answer"]
+        return response[0]["generated_text"]
     except Exception as e:
         return f"❌ Error processing image: {str(e)}"
 
@@ -121,15 +122,18 @@ img_interface = gr.Interface(
     fn=answer_question_from_image,
     inputs=[gr.Image(label="🖼️ Upload Image"), gr.Textbox(label="💬 Ask a Question")],
     outputs="text",
-    title="🖼️ AI Image Question Answering (NO OCR)"
+    title="🖼️ AI Image Question Answering"
 )
 
-# Define Gradio App
-app_ui = gr.TabbedInterface(
-    [doc_interface, img_interface], 
-    ["📄 Document QA", "🖼️ Image QA"]
-)
+# Mount Gradio Interfaces
+demo = gr.TabbedInterface([doc_interface, img_interface], ["📄 Document QA", "🖼️ Image QA"])
+app = gr.mount_gradio_app(app, demo, path="/")
+
+@app.get("/")
+def home():
+    return RedirectResponse(url="/")
 
-# Run Gradio UI separately
+# Run FastAPI + Gradio together
 if __name__ == "__main__":
-    app_ui.launch(share=True)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)