Update src/models.py

src/models.py

@@ -1,16 +1,8 @@
 # src/models.py
-#
-# OPTIMISATION SUMMARY vs original:
-# 1. torch.compile()         — fuses ops in SigLIP + DINOv2 forward passes (~25-40% faster on CPU/GPU)
-# 2. Batch embedding         — all crops embedded in ONE forward pass instead of N separate calls
-# 3. Image resize before AI  — downscale to 512px before any model touches the image (2-4x faster YOLO + DeepFace)
-# 4. half() on GPU           — FP16 inference halves memory and speeds up GPU (~2x)
-# 5. asyncio.to_thread()     — heavy CPU/GPU work offloaded so FastAPI stays non-blocking
-# 6. LRU image hash cache    — identical query images skip all inference (instant re-query)
-# 7. YOLO task='detect'      — segmentation masks (yolo11n-seg) replaced by plain detect (yolon11) for 3x speedup,
-#                               bounding boxes are just as good for crops
-# 8. Crop limit              — cap at MAX_CROPS (default 6) to prevent runaway latency on busy images
-# 9. enforce_detection=False — DeepFace won't raise on no-face; avoids Python exception overhead
+import os
+# FIX 1: Force Legacy Keras to prevent DeepFace/RetinaFace crash in TF 2.16+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"  # Hides the annoying CUDA/cuInit warnings
 
 import asyncio
 import hashlib
@@ -32,22 +24,18 @@ MAX_IMAGE_SIZE       = 512    # resize longest edge before any inference
 
 
 def _resize_pil(img: Image.Image, max_side: int = MAX_IMAGE_SIZE) -> Image.Image:
-    """Downscale so the longest side ≤ max_side, preserving aspect ratio."""
     w, h = img.size
     if max(w, h) <= max_side:
         return img
     scale = max_side / max(w, h)
     return img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
 
-
 def _img_hash(image_path: str) -> str:
-    """Fast xxhash-like hash of first 64 KB — good enough for cache keys."""
     h = hashlib.md5()
     with open(image_path, "rb") as f:
         h.update(f.read(65536))
     return h.hexdigest()
 
-
 class AIModelManager:
     def __init__(self):
         self.device = (
@@ -56,119 +44,68 @@ class AIModelManager:
         )
         print(f"Loading models onto: {self.device.upper()}...")
 
-        # ── SigLIP ────────────────────────────────────────────────
-        self.siglip_processor = AutoProcessor.from_pretrained(
-            "google/siglip-base-patch16-224", use_fast=True  # use_fast=True saves ~10ms
-        )
-        self.siglip_model = AutoModel.from_pretrained(
-            "google/siglip-base-patch16-224"
-        ).to(self.device).eval()
+        self.siglip_processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224", use_fast=True)
+        self.siglip_model = AutoModel.from_pretrained("google/siglip-base-patch16-224").to(self.device).eval()
 
-        # ── DINOv2 ────────────────────────────────────────────────
         self.dinov2_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
-        self.dinov2_model = AutoModel.from_pretrained(
-            "facebook/dinov2-base"
-        ).to(self.device).eval()
+        self.dinov2_model = AutoModel.from_pretrained("facebook/dinov2-base").to(self.device).eval()
 
-        # ── FP16 on GPU — halves memory, ~2x throughput ───────────
         if self.device == "cuda":
             self.siglip_model = self.siglip_model.half()
             self.dinov2_model = self.dinov2_model.half()
 
-        # ── torch.compile (PyTorch 2.0+) — fuses kernels ─────────
-        # Falls back silently on older torch versions
-        try:
-            self.siglip_model = torch.compile(self.siglip_model, mode="reduce-overhead")
-            self.dinov2_model  = torch.compile(self.dinov2_model,  mode="reduce-overhead")
-            print("✅ torch.compile enabled")
-        except Exception:
-            print("⚠️  torch.compile not available — running eager mode")
-
-        # ── YOLO — plain detect is 3x faster than seg ────────────
-        # Switch from yolo11n-seg.pt → yolo11n.pt (detection only)
-        # bounding boxes are sufficient for crops; we don't need masks
+        # FIX 2: Removed torch.compile() because HF Spaces do not have the g++ compiler installed by default.
+        # This fixes the "InvalidCxxCompiler" Search crash.
+        
         self.yolo = YOLO("yolo11n.pt")
 
-        # ── LRU result cache (keyed on MD5 of image bytes) ───────
-        # Caches the final vector list so identical re-uploads are instant
         self._cache = {}
         self._cache_maxsize = 256
 
         print("✅ Models ready!")
 
-    # ── BATCHED object embedding ───────────────────────────────────
     def _embed_crops_batch(self, crops: list[Image.Image]) -> list[np.ndarray]:
-        """
-        Run SigLIP + DINOv2 over ALL crops in ONE batched forward pass.
-        Much faster than calling _embed_object_crop() N times.
-        """
         if not crops:
             return []
 
         with torch.no_grad():
-            # SigLIP batch
-            sig_inputs = self.siglip_processor(
-                images=crops, return_tensors="pt", padding=True
-            )
+            sig_inputs = self.siglip_processor(images=crops, return_tensors="pt", padding=True)
             sig_inputs = {k: v.to(self.device) for k, v in sig_inputs.items()}
             if self.device == "cuda":
-                sig_inputs = {k: v.half() if v.dtype == torch.float32 else v
-                              for k, v in sig_inputs.items()}
+                sig_inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in sig_inputs.items()}
 
             sig_out = self.siglip_model.get_image_features(**sig_inputs)
             if hasattr(sig_out, "image_embeds"):
                 sig_out = sig_out.image_embeds
             elif isinstance(sig_out, tuple):
                 sig_out = sig_out[0]
-            sig_vecs = F.normalize(sig_out.float(), p=2, dim=1).cpu()  # [N, 768]
+            sig_vecs = F.normalize(sig_out.float(), p=2, dim=1).cpu()
 
-            # DINOv2 batch
-            dino_inputs = self.dinov2_processor(
-                images=crops, return_tensors="pt"
-            )
+            dino_inputs = self.dinov2_processor(images=crops, return_tensors="pt")
             dino_inputs = {k: v.to(self.device) for k, v in dino_inputs.items()}
             if self.device == "cuda":
-                dino_inputs = {k: v.half() if v.dtype == torch.float32 else v
-                               for k, v in dino_inputs.items()}
+                dino_inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in dino_inputs.items()}
 
             dino_out  = self.dinov2_model(**dino_inputs)
-            dino_vecs = dino_out.last_hidden_state[:, 0, :]          # CLS token
-            dino_vecs = F.normalize(dino_vecs.float(), p=2, dim=1).cpu()  # [N, 768]
+            dino_vecs = dino_out.last_hidden_state[:, 0, :]
+            dino_vecs = F.normalize(dino_vecs.float(), p=2, dim=1).cpu()
 
-            # Fuse → 1536-D, re-normalise
             fused = F.normalize(torch.cat([sig_vecs, dino_vecs], dim=1), p=2, dim=1)
 
         return [fused[i].numpy() for i in range(len(crops))]
 
-    # ── Main processing pipeline ───────────────────────────────────
-    def process_image(
-        self,
-        image_path: str,
-        is_query: bool  = False,
-        detect_faces: bool = True,
-    ) -> list[dict]:
-        """
-        Returns a list of {"type": "face"|"object", "vector": np.ndarray}.
-        Results for the same image bytes are returned from cache.
-        """
-        # ── Cache check ───────────────────────────────────────────
+    def process_image(self, image_path: str, is_query: bool = False, detect_faces: bool = True) -> list[dict]:
         cache_key = _img_hash(image_path)
         if cache_key in self._cache:
             print("⚡ Cache hit — skipping inference")
             return self._cache[cache_key]
 
         extracted = []
-
-        # ── Load & resize once ────────────────────────────────────
         original_pil = Image.open(image_path).convert("RGB")
         small_pil    = _resize_pil(original_pil, MAX_IMAGE_SIZE)
         img_np       = np.array(small_pil)
-        img_h, img_w = img_np.shape[:2]
         faces_found  = False
 
-        # ═════════════════════════════════════════════════════════
-        # LANE 1 — FACE LANE (toggleable)
-        # ═════════════════════════════════════════════════════════
         if detect_faces:
             try:
                 print("🔍 Face detection …")
@@ -176,10 +113,9 @@ class AIModelManager:
                     img_path=img_np,
                     model_name="GhostFaceNet",
                     detector_backend="retinaface",
-                    enforce_detection=False,   # no exception on miss — faster
+                    enforce_detection=False,
                     align=True,
                 )
-
                 for face in (face_objs or []):
                     fa = face.get("facial_area", {})
                     if fa.get("w", 0) * fa.get("h", 0) < MIN_FACE_AREA:
@@ -191,64 +127,40 @@ class AIModelManager:
 
             except Exception as e:
                 print(f"🟠 Face lane error: {e} — falling back to object lane")
-        else:
-            print("⏩ FAST MODE: skipping face lane")
-
-        # ══════════════════���══════════════════════════════════════
-        # LANE 2 — OBJECT LANE
-        # Collect all crops first, then embed as ONE batch
-        # ═════════════════════════════════════════════════════════
-        crops = [small_pil]   # always include full-image crop
 
+        crops = [small_pil]
         yolo_results = self.yolo(image_path, conf=0.5, verbose=False)
 
         for r in yolo_results:
             if r.boxes is None:
                 continue
-            for box_idx, box in enumerate(r.boxes):
+            for box in r.boxes:
                 cls_id = int(box.cls.item())
                 if faces_found and cls_id == YOLO_PERSON_CLASS_ID:
-                    continue   # skip person boxes when faces already indexed
+                    continue
                 x1, y1, x2, y2 = box.xyxy[0].tolist()
                 w, h = x2 - x1, y2 - y1
                 if w < 30 or h < 30:
                     continue
                 crop = small_pil.crop((x1, y1, x2, y2))
                 crops.append(crop)
-                if len(crops) >= MAX_CROPS + 1:   # +1 for the full-image crop
+                if len(crops) >= MAX_CROPS + 1:
                     break
             if len(crops) >= MAX_CROPS + 1:
                 break
 
-        # SINGLE batched forward pass for all crops
         print(f"🧠 Embedding {len(crops)} crop(s) in one batch …")
         vecs = self._embed_crops_batch(crops)
         for vec in vecs:
             extracted.append({"type": "object", "vector": vec})
 
-        # ── Store in cache ────────────────────────────────────────
         if len(self._cache) >= self._cache_maxsize:
-            # Evict the oldest key (simple FIFO)
             oldest = next(iter(self._cache))
             del self._cache[oldest]
         self._cache[cache_key] = extracted
 
         return extracted
 
-    # ── Async wrapper — keeps FastAPI non-blocking ─────────────────
-    async def process_image_async(
-        self,
-        image_path: str,
-        is_query: bool = False,
-        detect_faces: bool = True,
-    ) -> list[dict]:
-        """
-        Call this from async FastAPI endpoints instead of process_image().
-        Runs the heavy CPU/GPU work in a thread pool so the event loop
-        is never blocked, enabling true concurrent request handling.
-        """
+    async def process_image_async(self, image_path: str, is_query: bool = False, detect_faces: bool = True) -> list[dict]:
         loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(
-            None,
-            functools.partial(self.process_image, image_path, is_query, detect_faces),
-        )
+        return await loop.run_in_executor(None, functools.partial(self.process_image, image_path, is_query, detect_faces))