videopix commited on
Commit
fc19af7
·
verified ·
1 Parent(s): d02eefa

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -18
app.py CHANGED
@@ -1,7 +1,7 @@
1
  import io
2
  import asyncio
3
  import os
4
- from fastapi import FastAPI, File, UploadFile, Depends, Header, HTTPException
5
  from fastapi.responses import JSONResponse, HTMLResponse
6
  from PIL import Image
7
  import torch
@@ -14,18 +14,18 @@ app = FastAPI(title="Florence Image Caption API")
14
 
15
  device = "cuda" if torch.cuda.is_available() else "cpu"
16
 
17
- # Lazy-loaded model and processor
18
  processor = None
19
  model = None
20
  model_lock = asyncio.Lock()
21
 
22
- # ---------------------------------------------------
23
- # Token Authentication
24
- # ---------------------------------------------------
25
 
26
- API_TOKEN = os.getenv("img2caption") # secret token from your environment
27
 
28
- def check_token(auth: str):
 
 
 
29
  if auth is None or not auth.startswith("Bearer "):
30
  raise HTTPException(status_code=401, detail="Missing Authorization header")
31
 
@@ -33,8 +33,9 @@ def check_token(auth: str):
33
  if token != API_TOKEN:
34
  raise HTTPException(status_code=403, detail="Invalid token")
35
 
 
36
  # ---------------------------------------------------
37
- # Lazy Florence Model Load
38
  # ---------------------------------------------------
39
  async def load_model():
40
  global processor, model
@@ -76,7 +77,7 @@ def run_caption(image: Image.Image) -> str:
76
 
77
 
78
  # ---------------------------------------------------
79
- # Public Login Page
80
  # ---------------------------------------------------
81
  @app.get("/", response_class=HTMLResponse)
82
  def login_page():
@@ -85,7 +86,8 @@ def login_page():
85
  <html>
86
  <head><title>Login</title></head>
87
  <body style="font-family:Arial;max-width:500px;margin:40px auto;">
88
- <h2>Access Token Required</h2>
 
89
  <input id="token" type="password" style="width:100%;padding:10px;" placeholder="Enter token">
90
  <button onclick="login()" style="padding:10px;margin-top:10px;width:100%;">Continue</button>
91
 
@@ -93,7 +95,6 @@ def login_page():
93
  function login() {
94
  const t = document.getElementById("token").value;
95
  if (!t) return alert("Token required");
96
-
97
  sessionStorage.setItem("authToken", t);
98
  window.location.href = "/ui";
99
  }
@@ -105,12 +106,10 @@ function login() {
105
 
106
 
107
  # ---------------------------------------------------
108
- # Protected UI
109
  # ---------------------------------------------------
110
  @app.get("/ui", response_class=HTMLResponse)
111
- async def ui_page(authorization: str = Header(None)):
112
- check_token(authorization)
113
-
114
  return """
115
  <!DOCTYPE html>
116
  <html>
@@ -140,7 +139,6 @@ async def ui_page(authorization: str = Header(None)):
140
 
141
  <script>
142
  let token = sessionStorage.getItem("authToken");
143
-
144
  if (!token) {
145
  alert("No token found, please login again.");
146
  window.location.href = "/";
@@ -188,11 +186,11 @@ async def ui_page(authorization: str = Header(None)):
188
 
189
 
190
  # ---------------------------------------------------
191
- # Protected Caption API
192
  # ---------------------------------------------------
193
  @app.post("/img2caption")
194
  async def img2caption(file: UploadFile = File(...), authorization: str = Header(None)):
195
- check_token(authorization)
196
 
197
  try:
198
  async with model_lock:
 
1
  import io
2
  import asyncio
3
  import os
4
+ from fastapi import FastAPI, File, UploadFile, Header, HTTPException
5
  from fastapi.responses import JSONResponse, HTMLResponse
6
  from PIL import Image
7
  import torch
 
14
 
15
  device = "cuda" if torch.cuda.is_available() else "cpu"
16
 
 
17
  processor = None
18
  model = None
19
  model_lock = asyncio.Lock()
20
 
21
+ # -------- TOKEN from HF Space Secrets ----------
22
+ API_TOKEN = os.getenv("img2caption") # your secret token
 
23
 
 
24
 
25
+ # ---------------------------------------------------
26
+ # Verify Token for API only
27
+ # ---------------------------------------------------
28
+ def verify_token(auth: str | None):
29
  if auth is None or not auth.startswith("Bearer "):
30
  raise HTTPException(status_code=401, detail="Missing Authorization header")
31
 
 
33
  if token != API_TOKEN:
34
  raise HTTPException(status_code=403, detail="Invalid token")
35
 
36
+
37
  # ---------------------------------------------------
38
+ # Lazy Load Model
39
  # ---------------------------------------------------
40
  async def load_model():
41
  global processor, model
 
77
 
78
 
79
  # ---------------------------------------------------
80
+ # PUBLIC LOGIN PAGE
81
  # ---------------------------------------------------
82
  @app.get("/", response_class=HTMLResponse)
83
  def login_page():
 
86
  <html>
87
  <head><title>Login</title></head>
88
  <body style="font-family:Arial;max-width:500px;margin:40px auto;">
89
+
90
+ <h2>Enter Access Token</h2>
91
  <input id="token" type="password" style="width:100%;padding:10px;" placeholder="Enter token">
92
  <button onclick="login()" style="padding:10px;margin-top:10px;width:100%;">Continue</button>
93
 
 
95
  function login() {
96
  const t = document.getElementById("token").value;
97
  if (!t) return alert("Token required");
 
98
  sessionStorage.setItem("authToken", t);
99
  window.location.href = "/ui";
100
  }
 
106
 
107
 
108
  # ---------------------------------------------------
109
+ # PUBLIC UI PAGE (no token required)
110
  # ---------------------------------------------------
111
  @app.get("/ui", response_class=HTMLResponse)
112
+ def ui_page():
 
 
113
  return """
114
  <!DOCTYPE html>
115
  <html>
 
139
 
140
  <script>
141
  let token = sessionStorage.getItem("authToken");
 
142
  if (!token) {
143
  alert("No token found, please login again.");
144
  window.location.href = "/";
 
186
 
187
 
188
  # ---------------------------------------------------
189
+ # PROTECTED API ENDPOINT
190
  # ---------------------------------------------------
191
  @app.post("/img2caption")
192
  async def img2caption(file: UploadFile = File(...), authorization: str = Header(None)):
193
+ verify_token(authorization)
194
 
195
  try:
196
  async with model_lock: