fix: added

README.md

@@ -11,6 +11,8 @@ app_port: 7860
 
 FastAPI service for `Salesforce/blip-image-captioning-large`, packaged for Hugging Face Docker Spaces.
 
+Set a Hugging Face Space secret named `API_KEY` to protect the caption endpoint.
+
 ## Endpoints
 
 - `GET /health`
@@ -20,6 +22,7 @@ FastAPI service for `Salesforce/blip-image-captioning-large`, packaged for Huggi
 
 ```bash
 curl -X POST "http://localhost:7860/caption" \
+  -H "x-api-key: your-secret-key" \
   -F "image=@example.jpg" \
   -F "min_new_tokens=5" \
   -F "max_new_tokens=20"
@@ -29,5 +32,5 @@ curl -X POST "http://localhost:7860/caption" \
 
 ```bash
 docker build -t blip-large-api .
-docker run -p 7860:7860 blip-large-api
+docker run -p 7860:7860 -e API_KEY=your-secret-key blip-large-api
 ```

app/main.py

@@ -3,14 +3,17 @@ import os
 import time
 from contextlib import asynccontextmanager
 
-from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, Security, UploadFile
 from fastapi.responses import JSONResponse
+from fastapi.security import APIKeyHeader
 from PIL import Image, UnidentifiedImageError
 from transformers import AutoTokenizer, BlipForConditionalGeneration, BlipImageProcessor, BlipProcessor
 
 Image.MAX_IMAGE_PIXELS = None
 
 MODEL_ID = os.getenv("MODEL_ID", "Salesforce/blip-image-captioning-large")
+API_KEY = os.getenv("API_KEY")
+API_KEY_HEADER_NAME = os.getenv("API_KEY_HEADER_NAME", "x-api-key")
 USE_FAST_PROCESSOR = os.getenv("USE_FAST_PROCESSOR", "true").strip().lower() in {
     "1",
     "true",
@@ -21,6 +24,7 @@ MAX_IMAGE_SIZE = (128,128)
 
 processor: BlipProcessor | None = None
 model: BlipForConditionalGeneration | None = None
+api_key_header = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=False)
 
 
 def load_model() -> None:
@@ -33,6 +37,13 @@ def load_model() -> None:
         model = BlipForConditionalGeneration.from_pretrained(MODEL_ID)
 
 
+def verify_api_key(api_key: str | None = Security(api_key_header)) -> None:
+    if not API_KEY:
+        raise RuntimeError("API_KEY environment variable is required.")
+    if api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key.")
+
+
 def generate_caption(image_bytes: bytes, min_new_tokens: int, max_new_tokens: int) -> str:
     if processor is None or model is None:
         raise RuntimeError("Model is not loaded.")
@@ -85,6 +96,7 @@ async def health() -> dict[str, str]:
 
 @app.post("/caption")
 async def caption_image(
+    _: None = Security(verify_api_key),
     image: UploadFile = File(...),
     min_new_tokens: int = Form(5),
     max_new_tokens: int = Form(20),