Update appImage.py

appImage.py

@@ -1,27 +1,54 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 import easyocr
 from fastapi import FastAPI
-from fastapi.responses import RedirectResponse
+from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
 import tempfile
 import os
 from gtts import gTTS
 from fpdf import FPDF
 import datetime
+from PIL import Image
+import torch
 
 # Initialize components
 app = FastAPI()
 
-# Load models
-captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+# Load models - Using microsoft/git-large-coco
+try:
+    # Try loading the better model first
+    processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+    git_model = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco")
+    print("Successfully loaded microsoft/git-large-coco model")
+    USE_GIT = True
+except Exception as e:
+    print(f"Failed to load GIT model: {e}. Falling back to smaller model")
+    captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+    USE_GIT = False
+
+# Initialize EasyOCR
 reader = easyocr.Reader(['en', 'fr'])  # English and French OCR
 
+def generate_caption(image_path):
+    """Generate caption using the best available model"""
+    try:
+        if USE_GIT:
+            image = Image.open(image_path)
+            inputs = processor(images=image, return_tensors="pt")
+            outputs = git_model.generate(**inputs, max_length=50)
+            return processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        else:
+            result = captioner(image_path)
+            return result[0]['generated_text']
+    except Exception as e:
+        print(f"Caption generation error: {e}")
+        return "Could not generate caption"
+
 def analyze_image(image_path):
     """Process image with both captioning and OCR"""
     try:
         # Generate image caption
-        caption_result = captioner(image_path)
-        caption = caption_result[0]['generated_text']
+        caption = generate_caption(image_path)
         
         # Extract text with EasyOCR
         ocr_result = reader.readtext(image_path, detail=0)
@@ -167,4 +194,4 @@ app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")
 def redirect_to_interface():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")