Update app.py

app.py

@@ -152,8 +152,111 @@ app = gr.mount_gradio_app(app, demo, path="/")
 def home():
     return RedirectResponse(url="/")
 """
+import gradio as gr
+import numpy as np
+import fitz  # PyMuPDF
 import torch
-print("CUDA Available:", torch.cuda.is_available())
-print("Torch Device Count:", torch.cuda.device_count())
-print("Current Device:", torch.cuda.current_device() if torch.cuda.is_available() else "CPU")
-print("CUDA Device Name:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "None")
+from fastapi import FastAPI
+from transformers import pipeline
+from PIL import Image
+from starlette.responses import RedirectResponse
+from openpyxl import load_workbook
+from docx import Document
+from pptx import Presentation
+
+# ✅ Initialize FastAPI
+app = FastAPI()
+
+# ✅ Check if CUDA is Available (For Debugging)
+device = "cpu"
+print(f"✅ Running on: {device}")
+
+# ✅ Lazy Load Model Function (Loads Only When Needed)
+def get_qa_pipeline():
+    print("🔄 Loading QA Model on CPU...")
+    return pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)
+
+def get_image_captioning_pipeline():
+    print("🔄 Loading Image Captioning Model on CPU...")
+    return pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning", device=-1)
+
+# ✅ File Type Validation
+ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}
+
+def validate_file_type(file):
+    print(f"📂 Validating file: {file.name}")
+    ext = file.name.split(".")[-1].lower()
+    return None if ext in ALLOWED_EXTENSIONS else f"❌ Unsupported file format: {ext}"
+
+# ✅ Extract Text Functions (Optimized)
+def extract_text_from_pdf(file):
+    print("📄 Extracting text from PDF...")
+    with fitz.open(file.name) as doc:
+        return " ".join(page.get_text() for page in doc)
+
+def extract_text_from_docx(file):
+    print("📄 Extracting text from DOCX...")
+    doc = Document(file.name)
+    return " ".join(p.text for p in doc.paragraphs)
+
+def extract_text_from_pptx(file):
+    print("📄 Extracting text from PPTX...")
+    ppt = Presentation(file.name)
+    return " ".join(shape.text for slide in ppt.slides for shape in slide.shapes if hasattr(shape, "text"))
+
+def extract_text_from_excel(file):
+    print("📊 Extracting text from Excel...")
+    wb = load_workbook(file.name, data_only=True)
+    return " ".join(" ".join(str(cell) for cell in row if cell) for sheet in wb.worksheets for row in sheet.iter_rows(values_only=True))
+
+# ✅ Question Answering Function (Efficient Processing)
+async def answer_question(file, question: str):
+    print("🔍 Processing file for QA...")
+    
+    validation_error = validate_file_type(file)
+    if validation_error:
+        return validation_error
+
+    file_ext = file.name.split(".")[-1].lower()
+    text = ""
+
+    if file_ext == "pdf":
+        text = extract_text_from_pdf(file)
+    elif file_ext == "docx":
+        text = extract_text_from_docx(file)
+    elif file_ext == "pptx":
+        text = extract_text_from_pptx(file)
+    elif file_ext == "xlsx":
+        text = extract_text_from_excel(file)
+    
+    if not text.strip():
+        return "⚠️ No text extracted from the document."
+
+    print("✂️ Truncating text for faster processing...")
+    truncated_text = text[:1024]  # Reduce to 1024 characters for better speed
+
+    qa_pipeline = get_qa_pipeline()
+    response = qa_pipeline(f"Question: {question}\nContext: {truncated_text}")
+
+    return response[0]["generated_text"]
+
+# ✅ Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 📄 AI-Powered Document & Image QA")
+
+    with gr.Row():
+        file_input = gr.File(label="Upload Document")
+        question_input = gr.Textbox(label="Ask a Question", placeholder="What is this document about?")
+    
+    answer_output = gr.Textbox(label="Answer")
+    submit_btn = gr.Button("Get Answer")
+    
+    submit_btn.click(answer_question, inputs=[file_input, question_input], outputs=answer_output)
+
+# ✅ Mount Gradio with FastAPI
+app = gr.mount_gradio_app(app, demo, path="/demo")
+
+@app.get("/")
+def home():
+    return RedirectResponse(url="/demo")
+