Spaces:

videopix
/

image2caption

Running

App Files Files Community

videopix commited on Dec 2, 2025

Commit

735c9b7

verified ·

1 Parent(s): fc19af7

Update app.py

Browse files

Files changed (1) hide show

app.py +58 -83

app.py CHANGED Viewed

@@ -1,11 +1,14 @@
 import io
 import asyncio
-import os
-from fastapi import FastAPI, File, UploadFile, Header, HTTPException
 from fastapi.responses import JSONResponse, HTMLResponse
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
 # ---------------------------------------------------
 # FastAPI App
@@ -14,30 +17,16 @@ app = FastAPI(title="Florence Image Caption API")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 processor = None
 model = None
 model_lock = asyncio.Lock()
-# -------- TOKEN from HF Space Secrets ----------
-API_TOKEN = os.getenv("img2caption")   # your secret token
-# ---------------------------------------------------
-# Verify Token for API only
-# ---------------------------------------------------
-def verify_token(auth: str | None):
-    if auth is None or not auth.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Missing Authorization header")
-    token = auth.split("Bearer ")[-1].strip()
-    if token != API_TOKEN:
-        raise HTTPException(status_code=403, detail="Invalid token")
-# ---------------------------------------------------
-# Lazy Load Model
-# ---------------------------------------------------
 async def load_model():
     global processor, model
     if model is None:
@@ -52,6 +41,7 @@ async def load_model():
 def run_caption(image: Image.Image) -> str:
     inputs = processor(
         text="<MORE_DETAILED_CAPTION>",
         images=image,
@@ -77,39 +67,46 @@ def run_caption(image: Image.Image) -> str:
 # ---------------------------------------------------
-# PUBLIC LOGIN PAGE
 # ---------------------------------------------------
-@app.get("/", response_class=HTMLResponse)
-def login_page():
-    return """
-<!DOCTYPE html>
-<html>
-<head><title>Login</title></head>
-<body style="font-family:Arial;max-width:500px;margin:40px auto;">
-    <h2>Enter Access Token</h2>
-    <input id="token" type="password" style="width:100%;padding:10px;" placeholder="Enter token">
-    <button onclick="login()" style="padding:10px;margin-top:10px;width:100%;">Continue</button>
-<script>
-function login() {
-    const t = document.getElementById("token").value;
-    if (!t) return alert("Token required");
-    sessionStorage.setItem("authToken", t);
-    window.location.href = "/ui";
-}
-</script>
-</body>
-</html>
-"""
 # ---------------------------------------------------
-# PUBLIC UI PAGE (no token required)
 # ---------------------------------------------------
-@app.get("/ui", response_class=HTMLResponse)
-def ui_page():
     return """
 <!DOCTYPE html>
 <html>
@@ -118,11 +115,19 @@ def ui_page():
     <style>
         body { font-family: Arial; max-width: 650px; margin: 40px auto; }
         h2 { text-align: center; }
-        #preview { width: 100%; margin-top: 15px; display: none; border-radius: 8px; }
-        #captionBox { margin-top: 20px; padding: 15px; background: #eee; border-radius: 6px; display: none; }
-        button { padding: 12px; width: 100%; margin-top: 10px;
-                 background: #4A90E2; color: white; border: none;
-                 border-radius: 6px; cursor: pointer; font-size: 16px; }
         button:hover { background: #357ABD; }
     </style>
 </head>
@@ -138,12 +143,6 @@ def ui_page():
     <div id="captionBox"></div>
 <script>
-    let token = sessionStorage.getItem("authToken");
-    if (!token) {
-        alert("No token found, please login again.");
-        window.location.href = "/";
-    }
     const imageInput = document.getElementById("imageInput");
     const preview = document.getElementById("preview");
     const captionBox = document.getElementById("captionBox");
@@ -167,16 +166,15 @@ def ui_page():
         form.append("file", f);
         captionBox.style.display = "block";
-        captionBox.textContent = "Generating caption...";
         const res = await fetch("/img2caption", {
             method: "POST",
-            headers: { "Authorization": "Bearer " + token },
             body: form
         });
         const data = await res.json();
-        captionBox.textContent = data.caption || data.error;
     }
 </script>
@@ -185,33 +183,10 @@ def ui_page():
 """
-# ---------------------------------------------------
-# PROTECTED API ENDPOINT
-# ---------------------------------------------------
-@app.post("/img2caption")
-async def img2caption(file: UploadFile = File(...), authorization: str = Header(None)):
-    verify_token(authorization)
-    try:
-        async with model_lock:
-            await load_model()
-        data = await file.read()
-        image = Image.open(io.BytesIO(data)).convert("RGB")
-        caption = run_caption(image)
-        return {"caption": caption}
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-# ---------------------------------------------------
-# Local Run
-# ---------------------------------------------------
 def keep_alive():
     pass
 if __name__ == "__main__":
     import uvicorn
     print("🚀 Launching Fast img2caption API")

 import io
 import asyncio
+import threading
+import time
+from fastapi import FastAPI, File, UploadFile, Header
 from fastapi.responses import JSONResponse, HTMLResponse
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
+import requests
+import os
 # ---------------------------------------------------
 # FastAPI App
 device = "cuda" if torch.cuda.is_available() else "cpu"
+# Lazy load model on first request (prevents HF timeout)
 processor = None
 model = None
 model_lock = asyncio.Lock()
+# Hugging Face token stored in HF Secrets
+HF_TOKEN = os.getenv("img2caption")
 async def load_model():
+    """Load Florence model only when first needed."""
     global processor, model
     if model is None:
 def run_caption(image: Image.Image) -> str:
+    """Perform caption generation."""
     inputs = processor(
         text="<MORE_DETAILED_CAPTION>",
         images=image,
 # ---------------------------------------------------
+# API Endpoint (Token enforced only when app sends a token)
 # ---------------------------------------------------
+@app.post("/img2caption")
+async def img2caption(
+    file: UploadFile = File(...),
+    authorization: str = Header(None)
+):
+    # Apps must send token → enforce check
+    # UI sends no token → skip check → allow
+    if authorization is not None:
+        if not authorization.startswith("Bearer "):
+            return JSONResponse({"error": "Invalid token format"}, status_code=403)
+        token = authorization.replace("Bearer ", "").strip()
+        if token != HF_TOKEN:
+            return JSONResponse({"error": "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}
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
 # ---------------------------------------------------
+# Simple HTML UI (NO token required)
 # ---------------------------------------------------
+@app.get("/", response_class=HTMLResponse)
+def ui():
     return """
 <!DOCTYPE html>
 <html>
     <style>
         body { font-family: Arial; max-width: 650px; margin: 40px auto; }
         h2 { text-align: center; }
+        #preview {
+            width: 100%; margin-top: 15px; display: none;
+            border-radius: 8px;
+        }
+        #captionBox {
+            margin-top: 20px; padding: 15px;
+            background: #eee; border-radius: 6px; display: none;
+        }
+        button {
+            padding: 12px; width: 100%; margin-top: 10px;
+            background: #4A90E2; color: white; border: none;
+            border-radius: 6px; cursor: pointer; font-size: 16px;
+        }
         button:hover { background: #357ABD; }
     </style>
 </head>
     <div id="captionBox"></div>
 <script>
     const imageInput = document.getElementById("imageInput");
     const preview = document.getElementById("preview");
     const captionBox = document.getElementById("captionBox");
         form.append("file", f);
         captionBox.style.display = "block";
+        captionBox.innerHTML = "Generating caption...";
         const res = await fetch("/img2caption", {
             method: "POST",
             body: form
         });
         const data = await res.json();
+        captionBox.innerHTML = data.caption || data.error;
     }
 </script>
 """
 def keep_alive():
     pass
 if __name__ == "__main__":
     import uvicorn
     print("🚀 Launching Fast img2caption API")