Update appImage.py

appImage.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File
+"""from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import RedirectResponse, JSONResponse
 from transformers import AutoProcessor, AutoModelForCausalLM
 from PIL import Image
@@ -42,4 +42,26 @@ async def caption_from_frontend(file: UploadFile = File(...)):
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")"""
+# app_image_logic.py
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
+from PIL import Image
+import torch
+
+try:
+    processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco")
+    USE_GIT = True
+except:
+    captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+    USE_GIT = False
+
+def generate_caption(image_path):
+    if USE_GIT:
+        image = Image.open(image_path)
+        inputs = processor(images=image, return_tensors="pt")
+        outputs = 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']