Update app.py

app.py

@@ -1,11 +1,15 @@
+import os
+import secrets
 import spaces
 import torch
 import gradio as gr
 from typing import List, Union, Optional
+from fastapi import Header, HTTPException, status
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 
 MODEL_ID = "Dibachain/Diba-Embed"
+API_KEY = os.environ.get("MY_API_KEY", "")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f"Loading {MODEL_ID}...")
@@ -18,9 +22,26 @@ model = SentenceTransformer(
 model.max_seq_length = 512
 print("Model loaded successfully!")
 
-# 1. Put @spaces.GPU directly on the function bound to Gradio so ZeroGPU detects it at startup
+def verify_key(provided_key: Optional[str]):
+    if not API_KEY:
+        raise RuntimeError("MY_API_KEY secret is not configured in Space Settings.")
+    if not provided_key or not secrets.compare_digest(provided_key, API_KEY):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing API Key"
+        )
+
+# 1. ZeroGPU function protected with an internal key check so random visitors can't use the UI
 @spaces.GPU(duration=30)
-def compute_embeddings(texts: Union[str, List[str]], input_type: str = "document", normalize: bool = True):
+def compute_embeddings(
+    api_key: str,
+    texts: Union[str, List[str]],
+    input_type: str = "document",
+    normalize: bool = True
+):
+    if not API_KEY or not secrets.compare_digest(api_key or "", API_KEY):
+        raise gr.Error("Unauthorized: Invalid API Key")
+
     text_list = [texts] if isinstance(texts, str) else texts
     prompt_name = "query" if input_type == "query" else None
 
@@ -35,40 +56,47 @@ def compute_embeddings(texts: Union[str, List[str]], input_type: str = "document
         for i, emb in enumerate(embeddings)
     ]
 
-# 2. Wire compute_embeddings directly to gr.Interface (required by ZeroGPU scanner)
+# 2. Gradio Interface (api_name=False hides it from public Gradio API clients)
 demo = gr.Interface(
     fn=compute_embeddings,
     inputs=[
+        gr.Textbox(label="API Key", type="password", placeholder="Enter MY_API_KEY..."),
         gr.Textbox(label="Input Text", placeholder="Enter text to embed..."),
         gr.Radio(["document", "query"], value="document", label="Input Type"),
         gr.Checkbox(value=True, label="Normalize")
     ],
     outputs=gr.JSON(label="Output"),
-    title="Diba-Embed API",
-    description="POST to `/v1/embeddings` for standard OpenAI-compatible JSON responses."
+    title="Diba-Embed Private API",
+    api_name=False
 )
 
-# 3. Start Gradio non-blocking so it binds port 7860 once and registers ZeroGPU
-demo.launch(prevent_thread_lock=True)
+# 3. Launch Gradio with show_api=False to disable public /gradio_api docs
+demo.launch(prevent_thread_lock=True, show_api=False)
 
-# 4. Attach your custom /v1/embeddings route directly to Gradio's internal FastAPI server (demo.app)
+# 4. Protected FastAPI /v1/embeddings endpoint
 class EmbedRequest(BaseModel):
     input: Union[str, List[str]]
     input_type: Optional[str] = "document"
     normalize: Optional[bool] = True
 
-@demo.app.get("/health")
-def health_check():
-    return {"status": "online", "model": MODEL_ID}
-
 @demo.app.post("/v1/embeddings")
-def create_embeddings(req: EmbedRequest):
-    data = compute_embeddings(req.input, req.input_type or "document", req.normalize)
+def create_embeddings(
+    req: EmbedRequest,
+    authorization: Optional[str] = Header(default=None),
+    x_api_key: Optional[str] = Header(default=None)
+):
+    # Accept either "Authorization: Bearer <MY_API_KEY>" or "X-API-Key: <MY_API_KEY>"
+    token = x_api_key
+    if not token and authorization and authorization.lower().startswith("bearer "):
+        token = authorization[7:].strip()
+
+    verify_key(token)
+
+    data = compute_embeddings(token, req.input, req.input_type or "document", req.normalize)
     return {
         "object": "list",
         "model": MODEL_ID,
         "data": data
     }
 
-# Keep the main process alive
 demo.block_thread()