Update app.py

app.py

@@ -5,8 +5,9 @@ from tika import parser  # Apache Tika for document parsing
 import openpyxl
 from pptx import Presentation
 from PIL import Image
-from transformers import pipeline
+from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration
 import gradio as gr
+import torch
 import numpy as np
 
 # Initialize FastAPI
@@ -15,8 +16,13 @@ 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)
-image_captioning_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+
+# Load Image Captioning Model
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+model = model.to(dtype=torch.float16)  # Quantizing to FP16
 
 print("✅ Models loaded")
 
@@ -101,7 +107,9 @@ def answer_question_from_image(image, question):
             image = Image.fromarray(image)  # Convert to PIL Image
         
         print("🖼️ Generating caption for image...")
-        caption = image_captioning_pipeline(image)[0]['generated_text']
+        inputs = processor(images=image, return_tensors="pt", use_fast=True).to(dtype=torch.float16)
+        output = model.generate(**inputs)
+        caption = processor.decode(output[0], skip_special_tokens=True)
 
         print("🤖 Answering question based on caption...")
         response = doc_qa_pipeline(f"Question: {question}\nContext: {caption}")