Update app.py

Add FastAPI endpoint with API key auth

app.py

@@ -1,36 +1,115 @@
+import os
+import io
+import time
+import base64
+import secrets
+import logging
+from functools import lru_cache
+
 import gradio as gr
-from transformers import VisionEncoderDecoderModel, TrOCRProcessor
-import torch
+from fastapi import FastAPI, HTTPException, Security, Depends
+from fastapi.security.api_key import APIKeyHeader
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 from PIL import Image
+import torch
+from transformers import VisionEncoderDecoderModel, TrOCRProcessor
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ── Config ────────────────────────────────────────────────────────────────────
+API_KEY = os.environ.get("API_KEY", "changeme")
+DEVICE  = "cuda" if torch.cuda.is_available() else "cpu"
+MODELS  = [
+    'anuashok/ocr-captcha-v3',
+    'anuashok/ocr-captcha-v2',
+    'anuashok/ocr-captcha-v1',
+    'microsoft/trocr-base-printed'
+]
+DEFAULT_MODEL = MODELS[0]
+
+# ── API Key Auth ──────────────────────────────────────────────────────────────
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+def verify_key(key: str = Security(api_key_header)):
+    if not key or not secrets.compare_digest(key, API_KEY):
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+    return key
+
+# ── Model Cache ───────────────────────────────────────────────────────────────
+@lru_cache(maxsize=4)
+def load_model(model_id: str):
+    logger.info(f"Loading {model_id} on {DEVICE}...")
+    processor = TrOCRProcessor.from_pretrained(model_id)
+    model = VisionEncoderDecoderModel.from_pretrained(model_id).to(DEVICE)
+    model.eval()
+    logger.info(f"✅ {model_id} ready")
+    return processor, model
 
-def recognize_captcha(input, mdl):
-    
-    # Load model and processor
-    processor = TrOCRProcessor.from_pretrained(mdl)
-    model = VisionEncoderDecoderModel.from_pretrained(mdl)
-    
-    # Prepare image
-    pixel_values = processor(input, return_tensors="pt").pixel_values
-
-    # Generate text
-    generated_ids = model.generate(pixel_values)
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    
-    return generated_text
+# ── Inference ─────────────────────────────────────────────────────────────────
+def recognize_captcha(image, model_id):
+    processor, model = load_model(model_id)
+    pixel_values = processor(image, return_tensors="pt").pixel_values.to(DEVICE)
+    with torch.no_grad():
+        generated_ids = model.generate(pixel_values, max_new_tokens=32)
+    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
 iface = gr.Interface(
     fn=recognize_captcha,
-    inputs=[
-        gr.Image(),
-        gr.Dropdown(
-            ['anuashok/ocr-captcha-v3','anuashok/ocr-captcha-v2','anuashok/ocr-captcha-v1','microsoft/trocr-base-printed'], label='Model to use'
-        )
-    ],
-    outputs=['text'],
-    title  = "Character Sequence Recognition From Captcha Image",
-    description = "Using some TrOCR models found on the HF Hub to test/break tough text captchas. Will you have to train your own?",
-    article="Created by Neeraj with ❤️ !!!"
+    inputs=[gr.Image(), gr.Dropdown(MODELS, label="Model", value=DEFAULT_MODEL)],
+    outputs=["text"],
+    title="CAPTCHA Solver",
+    description="API available at /solve-captcha-base64"
 )
 
-iface.queue(max_size=10)
-iface.launch()
+# ── FastAPI ───────────────────────────────────────────────────────────────────
+app = gr.mount_gradio_app(
+    FastAPI(title="CAPTCHA Solver API"),
+    iface,
+    path="/"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ── Schemas ───────────────────────────────────────────────────────────────────
+class SolveRequest(BaseModel):
+    image_base64: str
+    model: str = DEFAULT_MODEL
+
+class SolveResponse(BaseModel):
+    success: bool
+    text: str = ""
+    processing_time: float = 0.0
+    model_used: str = ""
+    error: str = ""
+
+# ── Endpoints ─────────────────────────────────────────────────────────────────
+@app.get("/health")
+def health():
+    return {"status": "ok", "device": DEVICE, "quantized": False, "default_model": DEFAULT_MODEL}
+
+@app.post("/solve-captcha-base64", response_model=SolveResponse)
+def solve(req: SolveRequest, _: str = Depends(verify_key)):
+    start = time.time()
+    try:
+        raw = req.image_base64
+        if "," in raw:
+            raw = raw.split(",", 1)[1]
+
+        image = Image.open(io.BytesIO(base64.b64decode(raw))).convert("RGB")
+        text  = recognize_captcha(image, req.model)
+
+        elapsed = time.time() - start
+        logger.info(f"Solved '{text}' in {elapsed:.2f}s")
+        return SolveResponse(success=True, text=text.strip(), processing_time=elapsed, model_used=req.model)
+
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        return SolveResponse(success=False, error=str(e), processing_time=time.time() - start)