Update appImage.py

appImage.py

@@ -63,10 +63,7 @@ from fastapi import FastAPI
 from fastapi.responses import RedirectResponse, JSONResponse, FileResponse
 import os
 from PIL import Image
-from transformers import (
-    ViltProcessor, ViltForQuestionAnswering,
-    AutoProcessor, GitForCausalLM
-)
+from transformers import ViltProcessor, ViltForQuestionAnswering, AutoProcessor, AutoModelForCausalLM
 from gtts import gTTS
 import easyocr
 import torch
@@ -76,14 +73,22 @@ from io import BytesIO
 
 app = FastAPI()
 
-# Models
+# Initialize models with optimized settings
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
-caption_model = GitForCausalLM.from_pretrained("microsoft/git-large-coco")
-reader = easyocr.Reader(['en', 'fr'])
+
+# Load GIT model with performance optimizations
+git_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+git_model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/git-large-coco",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+
+reader = easyocr.Reader(['en', 'fr'], gpu=torch.cuda.is_available())
 
 def classify_question(question: str):
+    """Optimized question classification"""
     q = question.lower()
     if any(w in q for w in ["text", "say", "written", "read"]):
         return "ocr"
@@ -91,6 +96,17 @@ def classify_question(question: str):
         return "caption"
     return "vqa"
 
+@torch.inference_mode()
+def generate_caption(image):
+    """Optimized caption generation with GIT model"""
+    try:
+        inputs = git_processor(images=image, return_tensors="pt").to(git_model.device)
+        outputs = git_model.generate(**inputs, max_length=50)
+        return git_processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    except Exception as e:
+        print(f"Caption generation error: {e}")
+        return "Could not generate caption"
+
 def answer_question_from_image(image, question):
     if image is None or not question.strip():
         return "Please upload an image and ask a question.", None
@@ -103,17 +119,16 @@ def answer_question_from_image(image, question):
             answer = " ".join([entry[1] for entry in result]) or "No readable text found."
 
         elif mode == "caption":
-            image_tensor = caption_processor(images=image, return_tensors="pt").pixel_values
-            generated_ids = caption_model.generate(image_tensor, max_new_tokens=64)
-            answer = caption_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            answer = generate_caption(image)
 
-        else:
+        else:  # VQA mode
             inputs = vqa_processor(image, question, return_tensors="pt")
             with torch.no_grad():
                 outputs = vqa_model(**inputs)
             predicted_id = outputs.logits.argmax(-1).item()
             answer = vqa_model.config.id2label[predicted_id]
 
+        # Generate audio response
         tts = gTTS(text=answer)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
             tts.save(tmp.name)
@@ -124,4 +139,4 @@ def answer_question_from_image(image, question):
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/templates/home.html")
+    return RedirectResponse(url="/templates/home.html")