Update app.py

app.py

@@ -1,26 +1,30 @@
-import gradio as gr
+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
 from PIL import Image
 from transformers import pipeline
-import torch
+import gradio as gr
 import numpy as np
 
-# Load Optimized Hugging Face Models
+# Initialize FastAPI
+app = FastAPI()
+
 print("🔄 Loading models...")
 
+# Load Hugging Face Models
 qa_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)
 image_captioning_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=-1, use_fast=True)
 
 print("✅ Models loaded (Optimized for Speed)")
 
 # Allowed File Extensions
-ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}
+ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx", "jpg", "jpeg", "png"}
 
 def validate_file_type(file):
-    ext = file.name.split(".")[-1].lower()
+    ext = file.filename.split(".")[-1].lower()
     if ext not in ALLOWED_EXTENSIONS:
         return f"❌ Unsupported file format: {ext}"
     return None
@@ -48,8 +52,16 @@ def extract_text_from_excel(excel_bytes):
             text.append(" ".join(map(str, row)))
     return "\n".join(text) if text else "⚠️ No text found."
 
-# Function to process document and answer question
-def answer_question_from_document(file, question):
+# Function to process file (document or image) and answer question
+def answer_question(file, question: str):
+    if isinstance(file, np.ndarray):  
+        # Image processing
+        image = Image.fromarray(file)  
+        caption = image_captioning_pipeline(image)[0]['generated_text']
+        response = qa_pipeline(f"Question: {question}\nContext: {caption}")
+        return response[0]["generated_text"]
+    
+    # Document processing
     validation_error = validate_file_type(file)
     if validation_error:
         return validation_error
@@ -74,29 +86,23 @@ def answer_question_from_document(file, question):
 
     return response[0]["generated_text"]
 
-# Function to process image and answer question
-def answer_question_from_image(image, question):
-    if isinstance(image, np.ndarray):  
-        image = Image.fromarray(image)  
-    
-    caption = image_captioning_pipeline(image)[0]['generated_text']
-    response = qa_pipeline(f"Question: {question}\nContext: {caption}")
-
-    return response[0]["generated_text"]
-
-# Gradio Interface
+# Gradio Interface for both images & documents
 interface = gr.Interface(
-    fn=lambda file, image, question: (
-        answer_question_from_document(file, question) if file else answer_question_from_image(image, question)
-    ),
-    inputs=[
-        gr.File(label="📂 Upload Document (PDF, DOCX, PPTX, XLSX)", optional=True),
-        gr.Image(label="🖼️ Upload Image", optional=True),
-        gr.Textbox(label="💬 Ask a Question")
-    ],
+    fn=answer_question,
+    inputs=[gr.File(label="📂 Upload Document or Image"), gr.Textbox(label="💬 Ask a Question")],
     outputs="text",
-    title="📄 AI Document & Image Question Answering",
-    description="Upload a **document** (PDF, DOCX, PPTX, XLSX) or an **image**, then ask a question about its content."
+    title="📄🖼️ AI Document & Image Question Answering"
 )
 
-interface.launch()
+# Mount Gradio with FastAPI
+demo = interface
+app = gr.mount_gradio_app(app, demo, path="/")
+
+@app.get("/")
+def home():
+    return RedirectResponse(url="/")
+
+# Run FastAPI + Gradio together
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)