Upload 2 files

utils/cache_manager.py

@@ -1,8 +1,27 @@
 from functools import lru_cache
 from typing import Iterable, List, Tuple
 
+import torch
+from PIL import Image
+
 
 BATCH_SIZE = 50
+_CACHE_MODEL = None
+_CACHE_PROCESSOR = None
+
+
+def configure_cache(model, processor) -> None:
+    """Bind the shared model and processor for cached inference."""
+    global _CACHE_MODEL, _CACHE_PROCESSOR
+    _CACHE_MODEL = model
+    _CACHE_PROCESSOR = processor
+
+
+def preprocess_image(img_path: str) -> Image.Image:
+    img = Image.open(img_path)
+    img = img.convert("RGB")
+    img.thumbnail((448, 448))
+    return img
 
 
 def _ensure_tuple(labels: Iterable[str]) -> Tuple[str, ...]:
@@ -12,23 +31,19 @@ def _ensure_tuple(labels: Iterable[str]) -> Tuple[str, ...]:
 
 
 @lru_cache(maxsize=5)
-def cached_inference(image_path, labels, model, processor):
-    import torch
-    from PIL import Image
+def _cached_logits(image_path: str, label_tuple: Tuple[str, ...]) -> List[float]:
+    if _CACHE_MODEL is None or _CACHE_PROCESSOR is None:
+        raise RuntimeError("Cache manager not configured with model and processor.")
 
-    label_tuple: Tuple[str, ...] = _ensure_tuple(labels)
-
-    with Image.open(image_path).convert("RGB") as img:
-        tensor_image = img.copy()
-
-    device = next(model.parameters()).device
-    dtype = next(model.parameters()).dtype
+    device = next(_CACHE_MODEL.parameters()).device
+    dtype = next(_CACHE_MODEL.parameters()).dtype
+    image = preprocess_image(image_path)
     logits: List[float] = []
 
     with torch.no_grad():
         for start in range(0, len(label_tuple), BATCH_SIZE):
-            batch = label_tuple[start : start + BATCH_SIZE]
-            inputs = processor(images=tensor_image, text=list(batch), return_tensors="pt", padding=True)
+            batch = list(label_tuple[start : start + BATCH_SIZE])
+            inputs = _CACHE_PROCESSOR(images=image, text=batch, return_tensors="pt", padding=True)
 
             prepared = {}
             for key, value in inputs.items():
@@ -40,10 +55,14 @@ def cached_inference(image_path, labels, model, processor):
                 else:
                     prepared[key] = value
 
-            outputs = model(**prepared)
-            batch_logits = outputs.logits_per_image[0].detach().cpu().tolist()
-            logits.extend(batch_logits)
+            outputs = _CACHE_MODEL(**prepared)
+            logits.extend(outputs.logits_per_image[0].detach().cpu().tolist())
 
-    tensor_image.close()
+    image.close()
     scores = torch.softmax(torch.tensor(logits), dim=0).tolist()
     return scores
+
+
+def cached_inference(image_path: str, labels: Iterable[str]) -> List[float]:
+    label_tuple = _ensure_tuple(labels)
+    return _cached_logits(image_path, label_tuple)