feat: add model=ensemble route and in-memory SHA-256 LRU cache

api/main.py

@@ -9,6 +9,8 @@ FastAPI service providing:
 
 import io
 import os
+import hashlib
+from collections import OrderedDict
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Dict, List, Optional
@@ -20,14 +22,18 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
 from .schemas import HealthResponse, ModelsResponse, PredictionResponse, ModelInfo
-from .model_server import TamperingModelServer
+from .model_server import TamperingModelServer, predict_ensemble
 
 MODEL_INFO = {
-    "tiny":  {"params": "563K",  "description": "Fastest - edge/mobile tampering detection profile"},
-    "base":  {"params": "3.28M", "description": "Balanced accuracy, splicing detection, and localization"},
-    "large": {"params": "14.78M","description": "Highest forensic precision for inpainting & composite extortion"},
+    "tiny":     {"params": "563K",   "description": "Fastest - edge/mobile tampering detection profile (~4 ms)"},
+    "base":     {"params": "3.28M",  "description": "Standard Profile - balanced accuracy and localized inpainting"},
+    "large":    {"params": "13.94M", "description": "Deep Profile - recommended for vintage/retro AI & degraded media"},
+    "ensemble": {"params": "17.22M", "description": "SOTA Production Ensemble - Conservative Fusion of Base + Large with Heatmap Max-Gating"},
 }
 
+PREDICTION_CACHE: OrderedDict = OrderedDict()
+MAX_CACHE_SIZE = 256
+
 model_registry: Dict[str, TamperingModelServer] = {}
 DEFAULT_VARIANT = os.environ.get("PATCH_MFFT_VARIANT", "base")
 IMG_SIZE = int(os.environ.get("PATCH_MFFT_IMG_SIZE", "384"))
@@ -87,12 +93,17 @@ async def health():
 
 @app.get("/models", response_model=ModelsResponse)
 async def list_models():
+    default_model = "ensemble" if ("base" in model_registry and "large" in model_registry) else (
+        DEFAULT_VARIANT if DEFAULT_VARIANT in model_registry else (
+            list(model_registry.keys())[0] if model_registry else "base"
+        )
+    )
     return ModelsResponse(
-        default=DEFAULT_VARIANT if DEFAULT_VARIANT in model_registry else list(model_registry.keys())[0],
+        default=default_model,
         models=[
             ModelInfo(
                 id=v,
-                loaded=v in model_registry,
+                loaded=(v in model_registry) or (v == "ensemble" and "base" in model_registry and "large" in model_registry),
                 params=MODEL_INFO[v]["params"],
                 description=MODEL_INFO[v]["description"],
             )
@@ -104,13 +115,14 @@ async def list_models():
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(
     file: UploadFile = File(...),
-    model: str = "large",
+    model: str = "ensemble",
 ):
-    variant = (model or DEFAULT_VARIANT).lower()
-    if variant not in model_registry:
+    variant = (model or "ensemble").lower()
+    valid_variants = list(model_registry.keys()) + (["ensemble"] if "base" in model_registry and "large" in model_registry else [])
+    if variant not in valid_variants:
         raise HTTPException(
             status_code=404,
-            detail=f"Model variant '{variant}' not found. Available: {list(model_registry.keys())}",
+            detail=f"Model variant '{variant}' not found. Available: {valid_variants}",
         )
 
     if not file.content_type or not file.content_type.startswith("image/"):
@@ -120,12 +132,36 @@ async def predict(
     if len(content) > 25 * 1024 * 1024:
         raise HTTPException(status_code=400, detail="Image size exceeds 25 MB limit.")
 
+    sha256_hash = hashlib.sha256(content).hexdigest()
+    cache_key = f"{sha256_hash}_{variant}"
+
+    # Instant In-Memory LRU Cache Check (< 1 ms latency)
+    if cache_key in PREDICTION_CACHE:
+        cached_result = dict(PREDICTION_CACHE[cache_key])
+        cached_result["cached"] = True
+        PREDICTION_CACHE.move_to_end(cache_key)
+        return PredictionResponse(**cached_result)
+
     try:
         image = Image.open(io.BytesIO(content)).convert("RGB")
     except Exception:
         raise HTTPException(status_code=400, detail="Corrupted image file.")
 
-    server = model_registry[variant]
-    result = server.predict(image, filename=file.filename)
+    if variant == "ensemble":
+        result = predict_ensemble(
+            model_registry["base"],
+            model_registry["large"],
+            image,
+            filename=file.filename,
+            sha256_hash=sha256_hash,
+        )
+    else:
+        server = model_registry[variant]
+        result = server.predict(image, filename=file.filename, sha256_hash=sha256_hash)
+
+    # Save to high-speed LRU memory cache
+    PREDICTION_CACHE[cache_key] = result
+    if len(PREDICTION_CACHE) > MAX_CACHE_SIZE:
+        PREDICTION_CACHE.popitem(last=False)
 
     return PredictionResponse(**result)