modifief app.py for private repo

app.py

@@ -12,8 +12,9 @@ import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 
-from fastapi import FastAPI, UploadFile, File, HTTPException, Request, Response
+from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security
 from fastapi.responses import RedirectResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket
 from bson import ObjectId
@@ -22,15 +23,11 @@ import uvicorn
 import gradio as gr
 from gradio import mount_gradio_app
 
-# -------------------------------------------------
-# Logging
-# -------------------------------------------------
+# --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# -------------------------------------------------
-# Paths
-# -------------------------------------------------
+# --------------------- Paths -----------------------
 REPO_ID = "HariLogicgo/face_swap_models"
 BASE_DIR = "./workspace"
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
@@ -41,17 +38,21 @@ os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULT_DIR, exist_ok=True)
 os.makedirs(MODELS_DIR, exist_ok=True)
 
-# -------------------------------------------------
-# Download models
-# -------------------------------------------------
+# --------------------- Secrets ---------------------
+HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face private repo token
+API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")  # Bearer token for API
+
+# --------------------- Download Models ---------------------
 def download_models():
-    logger.info("Downloading models...")
+    logger.info("Downloading models from private HF repo...")
     inswapper_path = hf_hub_download(
         repo_id=REPO_ID,
         filename="models/inswapper_128.onnx",
         repo_type="model",
-        local_dir=MODELS_DIR
+        local_dir=MODELS_DIR,
+        token=HF_TOKEN
     )
+
     buffalo_files = [
         "1k3d68.onnx",
         "2d106det.onnx",
@@ -64,24 +65,21 @@ def download_models():
             repo_id=REPO_ID,
             filename=f"models/buffalo_l/{f}",
             repo_type="model",
-            local_dir=MODELS_DIR
+            local_dir=MODELS_DIR,
+            token=HF_TOKEN
         )
     logger.info("Models downloaded successfully")
     return inswapper_path
 
 inswapper_path = download_models()
 
-# -------------------------------------------------
-# Face Analysis + Swapper
-# -------------------------------------------------
+# --------------------- Face Analysis + Swapper ---------------------
 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
 face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
 
-# -------------------------------------------------
-# CodeFormer setup
-# -------------------------------------------------
+# --------------------- CodeFormer ---------------------
 CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
 
 def ensure_codeformer():
@@ -94,9 +92,7 @@ def ensure_codeformer():
 
 ensure_codeformer()
 
-# -------------------------------------------------
-# MongoDB + GridFS setup
-# -------------------------------------------------
+# --------------------- MongoDB ---------------------
 MONGODB_URL = os.getenv(
     "MONGODB_URL",
     "mongodb+srv://harilogicgo_db_user:[email protected]/?retryWrites=true&w=majority&appName=Cluster0"
@@ -106,9 +102,7 @@ client: AsyncIOMotorClient = None
 database = None
 fs_bucket: AsyncIOMotorGridFSBucket = None
 
-# -------------------------------------------------
-# FastAPI App
-# -------------------------------------------------
+# --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
 
 @fastapi_app.on_event("startup")
@@ -127,39 +121,17 @@ async def shutdown_db():
         client.close()
         logger.info("MongoDB connection closed")
 
-# -------------------------------------------------
-# 🔐 Global Auth Middleware
-# -------------------------------------------------
-PUBLIC_PATHS = ["/health"]
-
-API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")
-if not API_SECRET_TOKEN:
-    raise RuntimeError("❌ API_SECRET_TOKEN not set in environment!")
-
-@fastapi_app.middleware("http")
-async def auth_middleware(request: Request, call_next):
-    path = request.url.path
-    if path in PUBLIC_PATHS:  # exact match
-        return await call_next(request)
-
-    auth_header = request.headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        return Response(content="Missing Bearer token", status_code=401)
-
-    token = auth_header.split("Bearer ")[1]
-    if token != API_SECRET_TOKEN:
-        return Response(content="Invalid token", status_code=403)
+# --------------------- Auth ---------------------
+security = HTTPBearer()
 
-    return await call_next(request)
+def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
+    if credentials.credentials != API_SECRET_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid or missing token")
+    return credentials.credentials
 
-# -------------------------------------------------
-# Lock for face swap
-# -------------------------------------------------
+# --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
-# -------------------------------------------------
-# Pipeline Function
-# -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
     try:
         with swap_lock:
@@ -201,9 +173,7 @@ def face_swap_and_enhance(src_img, tgt_img):
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
-# -------------------------------------------------
-# Gradio UI
-# -------------------------------------------------
+# --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")
 
@@ -222,24 +192,22 @@ with gr.Blocks() as demo:
 
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
 
-# -------------------------------------------------
-# API Endpoints
-# -------------------------------------------------
+# --------------------- API Endpoints ---------------------
 @fastapi_app.get("/")
 def root():
     return RedirectResponse("/gradio")
 
-@fastapi_app.get("/health")
+@fastapi_app.get("/health", dependencies=[Depends(verify_token)])
 async def health():
     return {"status": "healthy"}
 
-@fastapi_app.post("/source")
+@fastapi_app.post("/source", dependencies=[Depends(verify_token)])
 async def upload_source(image: UploadFile = File(...)):
     contents = await image.read()
     file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "source"})
     return {"source_id": str(file_id)}
 
-@fastapi_app.get("/targets")
+@fastapi_app.get("/targets", dependencies=[Depends(verify_token)])
 async def list_targets():
     files = []
     async for file in database.fs.files.find({"metadata.type": "target", "metadata.predefined": True}):
@@ -249,7 +217,7 @@ async def list_targets():
         })
     return {"targets": files}
 
-@fastapi_app.post("/target")
+@fastapi_app.post("/target", dependencies=[Depends(verify_token)])
 async def upload_target(image: UploadFile = File(...)):
     contents = await image.read()
     file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "target"})
@@ -259,7 +227,7 @@ class FaceSwapRequest(BaseModel):
     source_id: str
     target_id: str
 
-@fastapi_app.post("/faceswap")
+@fastapi_app.post("/faceswap", dependencies=[Depends(verify_token)])
 async def perform_faceswap(request: FaceSwapRequest):
     try:
         src_stream = await fs_bucket.open_download_stream(ObjectId(request.source_id))
@@ -290,7 +258,7 @@ async def perform_faceswap(request: FaceSwapRequest):
     result_id = await fs_bucket.upload_from_stream("enhanced.png", final_bytes, metadata={"type": "result"})
     return {"result_id": str(result_id)}
 
-@fastapi_app.get("/download/{result_id}")
+@fastapi_app.get("/download/{result_id}", dependencies=[Depends(verify_token)])
 async def download_result(result_id: str):
     try:
         stream = await fs_bucket.open_download_stream(ObjectId(result_id))
@@ -304,10 +272,8 @@ async def download_result(result_id: str):
         headers={"Content-Disposition": "attachment; filename=result.png"}
     )
 
-# -------------------------------------------------
-# Mount Gradio (last, without overwriting app)
-# -------------------------------------------------
-mount_gradio_app(fastapi_app, demo, path="/gradio")
+# --------------------- Mount Gradio ---------------------
+fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 
 if __name__ == "__main__":
     uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)