Update app_working_api.py

app_working_api.py

@@ -2,12 +2,13 @@ import io
 import asyncio
 import threading
 import time
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi import FastAPI, File, UploadFile, Header
+from fastapi.responses import JSONResponse, HTMLResponse, PlainTextResponse
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
 import requests
+import os
 
 # ---------------------------------------------------
 # FastAPI App
@@ -21,9 +22,10 @@ processor = None
 model = None
 model_lock = asyncio.Lock()
 
+# Hugging Face token stored in Space secrets
+HF_TOKEN = os.getenv("img2caption")
 
 async def load_model():
-    """Load Florence model only when first needed."""
     global processor, model
 
     if model is None:
@@ -38,7 +40,6 @@ async def load_model():
 
 
 def run_caption(image: Image.Image) -> str:
-    """Perform caption generation."""
     inputs = processor(
         text="<MORE_DETAILED_CAPTION>",
         images=image,
@@ -64,30 +65,41 @@ def run_caption(image: Image.Image) -> str:
 
 
 # ---------------------------------------------------
-# API Endpoint
+# API Endpoint (Protected only if token is sent)
 # ---------------------------------------------------
-@app.post("/img2caption")
-async def img2caption(file: UploadFile = File(...)):
+@app.post("/img2caption", response_class=PlainTextResponse)
+async def img2caption(
+    file: UploadFile = File(...),
+    authorization: str = Header(None)
+):
+
+    # If app sends a token → validate it
+    if authorization is not None:
+        if not authorization.startswith("Bearer "):
+            return PlainTextResponse("Invalid token format", status_code=403)
+
+        token = authorization.replace("Bearer ", "").strip()
+        if token != HF_TOKEN:
+            return PlainTextResponse("Invalid token", status_code=403)
+
     try:
-        # Ensure model is loaded
         async with model_lock:
             await load_model()
 
-        # Read and convert image
         data = await file.read()
         image = Image.open(io.BytesIO(data)).convert("RGB")
 
-        # Caption
         caption = run_caption(image)
 
-        return {"caption": caption}
+        # Return ONLY the caption string, no JSON
+        return caption
 
     except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
+        return PlainTextResponse(f"Error: {str(e)}", status_code=500)
 
 
 # ---------------------------------------------------
-# Simple HTML UI
+# Simple HTML UI (no token required)
 # ---------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 def ui():
@@ -146,8 +158,8 @@ def ui():
             method: "POST",
             body: form
         });
-        const data = await res.json();
-        captionBox.innerHTML = data.caption || data.error;
+        const text = await res.text();
+        captionBox.innerHTML = text;
     }
 </script>
 </body>
@@ -158,6 +170,7 @@ def ui():
 def keep_alive():
     pass
 
+
 if __name__ == "__main__":
     import uvicorn
     print("🚀 Launching Fast img2caption API")