Update app.py

app.py

@@ -1,3 +1,6 @@
+from fastapi import FastAPI, File, UploadFile
+import requests
+
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 import gradio as gr
@@ -17,6 +20,9 @@ import concurrent.futures
 # Load environment variables from .env file
 load_dotenv()
 
+app = FastAPI()
+
+
 # Salesforce credentials
 SF_USERNAME = os.getenv('SF_USERNAME')
 SF_PASSWORD = os.getenv('SF_PASSWORD')
@@ -36,6 +42,22 @@ model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+
+# FastAPI endpoint to handle image upload and caption generation
+@app.post("/predict/")
+async def predict(image: UploadFile = File(...)):
+    try:
+        # Read the image from the request
+        image_bytes = await image.read()
+        image = Image.open(BytesIO(image_bytes))
+        
+        # Generate caption from the image
+        caption = generate_captions_from_image(image)
+        return {"caption": caption}
+    except Exception as e:
+        return {"error": str(e)}
+
+
 # Inference function to generate captions dynamically based on image content
 def generate_captions_from_image(image):
     if image.mode != "RGB":
@@ -267,4 +289,6 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+