Update app.py

app.py

@@ -5,6 +5,9 @@ import base64
 import secrets
 import logging
 import numpy as np
+from hashlib import md5
+from collections import OrderedDict
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, HTTPException, Security, Depends
 from fastapi.security.api_key import APIKeyHeader
@@ -27,6 +30,47 @@ MODEL_PATH = "./model/model.onnx"
 
 ORT_INTRA_THREADS = int(os.getenv("ORT_INTRA_THREADS", "1"))
 ORT_INTER_THREADS = int(os.getenv("ORT_INTER_THREADS", "1"))
+CACHE_MAX_SIZE    = int(os.getenv("CACHE_MAX_SIZE", "500"))
+
+# ── Torch global optimizations ────────────────────────────────────────────────
+torch.set_grad_enabled(False)       # no autograd overhead on tensor ops
+torch.set_num_threads(1)            # don't compete with ORT threads
+torch.set_num_interop_threads(1)    # no inter-op parallelism from torch
+
+# ── LRU Cache ─────────────────────────────────────────────────────────────────
+class LRUCache:
+    def __init__(self, max_size: int):
+        self.max_size = max_size
+        self._cache: OrderedDict = OrderedDict()
+        self.hits   = 0
+        self.misses = 0
+
+    def get(self, key: str):
+        if key not in self._cache:
+            self.misses += 1
+            return None
+        self._cache.move_to_end(key)
+        self.hits += 1
+        return self._cache[key]
+
+    def set(self, key: str, value: str):
+        if key in self._cache:
+            self._cache.move_to_end(key)
+        else:
+            if len(self._cache) >= self.max_size:
+                self._cache.popitem(last=False)
+        self._cache[key] = value
+
+    @property
+    def size(self):
+        return len(self._cache)
+
+    @property
+    def hit_rate(self):
+        total = self.hits + self.misses
+        return round(self.hits / total * 100, 1) if total else 0.0
+
+cache = LRUCache(max_size=CACHE_MAX_SIZE)
 
 # ── API Key Auth ──────────────────────────────────────────────────────────────
 api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
@@ -46,28 +90,66 @@ opts.inter_op_num_threads     = ORT_INTER_THREADS
 opts.execution_mode           = rt.ExecutionMode.ORT_SEQUENTIAL
 opts.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
 opts.optimized_model_filepath = MODEL_PATH + ".opt"
+opts.enable_mem_pattern       = True   # reuse memory allocations across runs
+opts.enable_cpu_mem_arena     = True   # pre-allocate memory pool for ORT
 
 session = rt.InferenceSession(
     MODEL_PATH,
     sess_options=opts,
     providers=["CPUExecutionProvider"]
 )
-input_name = session.get_inputs()[0].name
+input_name = session.get_inputs()[0].name   # cache once at startup
 logger.info("✅ ONNX model ready")
 
 # ── Inference ─────────────────────────────────────────────────────────────────
+def preprocess(image: Image.Image) -> np.ndarray:
+    """Normalize and transpose in one contiguous block — no extra memory copies."""
+    image = image.convert("RGB").resize(IMG_SIZE, Image.BILINEAR)  # faster than BICUBIC
+    x = np.ascontiguousarray(image, dtype=np.float32)
+    x = (x / 255.0 - 0.5) / 0.5
+    return x.transpose(2, 0, 1)[np.newaxis, :]   # [1, 3, H, W]
+
+
 def solve_image(image: Image.Image) -> str:
-    image = image.convert("RGB").resize(IMG_SIZE, Image.BICUBIC)
-    x = np.array(image, dtype=np.float32) / 255.0
-    x = (x - 0.5) / 0.5
-    x = x.transpose(2, 0, 1)[np.newaxis, :]        # [1, 3, H, W]
-    logits = session.run(None, {input_name: x})[0]  # numpy output from ONNX
-    # tokenizer uses torch Tensor ops internally (max, slicing, tolist)
-    # so convert numpy → torch only at this boundary, nowhere else
-    probs = torch.tensor(logits).softmax(-1)
+    x      = preprocess(image)
+    logits = session.run(None, {input_name: x})[0]
+    # tokenizer uses torch Tensor ops internally — convert only at this boundary
+    probs  = torch.tensor(logits).softmax(-1)
     preds, _ = tokenizer.decode(probs)
     return preds[0]
 
+
+def solve_with_cache(raw_b64: str, image: Image.Image) -> tuple[str, bool]:
+    key    = md5(raw_b64.encode()).hexdigest()
+    cached = cache.get(key)
+    if cached is not None:
+        logger.info(f"⚡ Cache hit — '{cached}'")
+        return cached, True
+    text = solve_image(image).strip()[:5]
+    cache.set(key, text)
+    return text, False
+
+
+# ── Warmup ────────────────────────────────────────────────────────────────────
+def warmup():
+    """
+    Run 3 dummy inferences at startup so ORT JIT-compiles the graph
+    before any real request arrives. Every worker runs this independently.
+    """
+    logger.info("Warming up model...")
+    dummy = Image.new("RGB", IMG_SIZE, color=(128, 128, 128))
+    for _ in range(3):
+        solve_image(dummy)
+    logger.info("✅ Warmup complete — model is hot")
+
+
+# ── Lifespan ──────────────────────────────────────────────────────────────────
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    warmup()
+    yield
+
+
 # ── Schemas ───────────────────────────────────────────────────────────────────
 class SolveRequest(BaseModel):
     image_base64: str
@@ -76,10 +158,12 @@ class SolveResponse(BaseModel):
     success: bool
     text: str = ""
     processing_time: float = 0.0
+    cached: bool = False
     error: str = ""
 
+
 # ── FastAPI ───────────────────────────────────────────────────────────────────
-app = FastAPI(title="CAPTCHA Solver API")
+app = FastAPI(title="CAPTCHA Solver API", lifespan=lifespan)
 
 app.add_middleware(
     CORSMiddleware,
@@ -88,6 +172,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+
 # ── Endpoints ─────────────────────────────────────────────────────────────────
 @app.get("/health")
 def health():
@@ -98,8 +183,16 @@ def health():
         "quantized":     True,
         "workers":       os.getenv("WEB_CONCURRENCY", "1"),
         "intra_threads": ORT_INTRA_THREADS,
+        "cache": {
+            "size":     cache.size,
+            "max_size": cache.max_size,
+            "hits":     cache.hits,
+            "misses":   cache.misses,
+            "hit_rate": f"{cache.hit_rate}%",
+        },
     }
 
+
 @app.post("/solve-captcha-base64", response_model=SolveResponse)
 def solve(req: SolveRequest, _: str = Depends(verify_key)):
     start = time.time()
@@ -108,14 +201,18 @@ def solve(req: SolveRequest, _: str = Depends(verify_key)):
         if "," in raw:
             raw = raw.split(",", 1)[1]
 
-        image = Image.open(io.BytesIO(base64.b64decode(raw))).convert("RGB")
-        text  = solve_image(image)
-        text  = text.strip()[:5]
+        image = Image.open(io.BytesIO(base64.b64decode(raw)))
+        text, hit = solve_with_cache(raw, image)
 
         elapsed = time.time() - start
-        logger.info(f"✅ Solved: '{text}' in {elapsed:.3f}s")
+        logger.info(f"✅ Solved: '{text}' in {elapsed:.3f}s (cached={hit})")
 
-        return SolveResponse(success=True, text=text, processing_time=elapsed)
+        return SolveResponse(
+            success=True,
+            text=text,
+            processing_time=elapsed,
+            cached=hit,
+        )
 
     except Exception as e:
         logger.error(f"Error: {e}")