Spaces:
Sleeping
Sleeping
| """ | |
| backbone_manager.py | |
| ───────────────────────────────────────────────────────────────────── | |
| Group 07 — I-ILIA-014 MIR Project | |
| Phase 3: Centralised, multi-backbone resource manager. | |
| Replaces the global FAISS / embedder singletons in ``app.py`` with a | |
| thread-safe cache that: | |
| • lazy-loads embedders and FAISS indices on first use, | |
| • caches them as long as RAM allows, | |
| • evicts least-recently-used entries when the cache exceeds | |
| ``MIR_MAX_BACKBONES`` (env var, default 1). | |
| This makes the live backbone selector in the UI cheap to swap on a | |
| laptop (one backbone resident at a time) while still letting HF Spaces | |
| CPU-Basic (16 GB RAM) keep all four resident for instant switching. | |
| Public API: | |
| from backbone_manager import manager | |
| embedder = manager.get_embedder("clip") | |
| image_idx = manager.get_image_index("clip") # FAISS | |
| caption_idx = manager.get_caption_index("clip") # FAISS | |
| img_files = manager.get_image_files("clip") # list[str] | |
| captions = manager.get_captions("clip") # list[str] | |
| manager.list_backbones() → list of {"key", "backbone_id", | |
| "display_name", "loaded": bool} | |
| manager.is_known("clip") → bool | |
| ───────────────────────────────────────────────────────────────────── | |
| """ | |
| from __future__ import annotations | |
| import gc | |
| import logging | |
| import os | |
| import threading | |
| from collections import OrderedDict | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import h5py | |
| import numpy as np | |
| import faiss # type: ignore | |
| log = logging.getLogger("backbone_manager") | |
| # ───────────────────────────────────────────────────────────────────── | |
| # Registry — single source of truth, must match embedders/__init__.py | |
| # ───────────────────────────────────────────────────────────────────── | |
| class BackboneSpec: | |
| key: str | |
| backbone_id: str | |
| display_name: str | |
| BACKBONES: dict[str, BackboneSpec] = { | |
| "clip": BackboneSpec("clip", "clip_openai_b32", | |
| "CLIP ViT-B/32 (OpenAI)"), | |
| "openclip_b32": BackboneSpec("openclip_b32", "openclip_laion2b_vitb32", | |
| "OpenCLIP ViT-B/32 (LAION-2B)"), | |
| "openclip_l14": BackboneSpec("openclip_l14", "openclip_laion2b_vitl14", | |
| "OpenCLIP ViT-L/14 (LAION-2B)"), | |
| "blip": BackboneSpec("blip", "blip_itm_base", | |
| "BLIP ITM-base (COCO)"), | |
| } | |
| DEFAULT_BACKBONE = os.environ.get("MIR_DEFAULT_BACKBONE", "clip") | |
| MAX_BACKBONES = max(1, int(os.environ.get("MIR_MAX_BACKBONES", "1"))) | |
| H5_DIR = Path(os.environ.get("MIR_H5_DIR", "descriptors/CLIP_Flickr8k")) | |
| # ───────────────────────────────────────────────────────────────────── | |
| # Cached payload (embedder + FAISS indices + identifier lists) — one | |
| # per backbone key. We bundle them together so the LRU evicts as a | |
| # unit, since they're useless without each other. | |
| # ───────────────────────────────────────────────────────────────────── | |
| class _BackboneAssets: | |
| embedder: object # whatever get_embedder returns | |
| image_index: faiss.Index | |
| caption_index: faiss.Index | |
| image_files: list[str] | |
| captions: list[str] | |
| caption_filenames: list[str] # filenames corresponding to captions built | |
| embedding_dim: int | |
| # ───────────────────────────────────────────────────────────────────── | |
| # H5 loader (mirrors compare_backbones.py — same defensive key probing) | |
| # ───────────────────────────────────────────────────────────────────── | |
| def _load_h5(path: Path) -> tuple[np.ndarray, list[str]]: | |
| if not path.exists(): | |
| raise FileNotFoundError(path) | |
| with h5py.File(path, "r") as f: | |
| keys = list(f.keys()) | |
| emb_key = next((k for k in ("embeddings", "features", "vectors") if k in keys), None) | |
| if emb_key is None: | |
| raise KeyError(f"No embeddings dataset in {path} (keys={keys})") | |
| embeddings = f[emb_key][:] | |
| id_key = next((k for k in ("filenames", "captions", "ids", "labels") if k in keys), None) | |
| if id_key is None: | |
| raise KeyError(f"No identifier dataset in {path} (keys={keys})") | |
| raw = f[id_key][:] | |
| identifiers = [s.decode("utf-8") if isinstance(s, bytes) else str(s) for s in raw] | |
| return embeddings, identifiers | |
| def _build_faiss_index(vectors: np.ndarray) -> faiss.IndexFlatIP: | |
| vectors = vectors.astype(np.float32, copy=False) | |
| faiss.normalize_L2(vectors) | |
| idx = faiss.IndexFlatIP(vectors.shape[1]) | |
| idx.add(vectors) | |
| return idx | |
| # ───────────────────────────────────────────────────────────────────── | |
| # Manager | |
| # ───────────────────────────────────────────────────────────────────── | |
| class BackboneManager: | |
| """Thread-safe LRU cache for VLM backbones + their FAISS indices.""" | |
| def __init__(self, max_backbones: int = MAX_BACKBONES, | |
| h5_dir: Path = H5_DIR) -> None: | |
| self._max = max_backbones | |
| self._h5_dir = h5_dir | |
| self._cache: "OrderedDict[str, _BackboneAssets]" = OrderedDict() | |
| self._lock = threading.RLock() | |
| log.info("BackboneManager initialised · MAX_BACKBONES=%d · h5_dir=%s", | |
| max_backbones, h5_dir) | |
| # ───── Public ───── | |
| def is_known(self, key: str) -> bool: | |
| return key in BACKBONES | |
| def list_backbones(self) -> list[dict]: | |
| """For the /backbones API endpoint that powers the dropdown.""" | |
| with self._lock: | |
| return [ | |
| { | |
| "key": spec.key, | |
| "backbone_id": spec.backbone_id, | |
| "display_name": spec.display_name, | |
| "loaded": spec.key in self._cache, | |
| "active": False, # the route fills this from session | |
| } | |
| for spec in BACKBONES.values() | |
| ] | |
| def get_caption_filenames(self, key: str) -> list[str]: | |
| return self._get_assets(key).caption_filenames | |
| def get_embedder(self, key: str): | |
| return self._get_assets(key).embedder | |
| def get_image_index(self, key: str) -> faiss.Index: | |
| return self._get_assets(key).image_index | |
| def get_caption_index(self, key: str) -> faiss.Index: | |
| return self._get_assets(key).caption_index | |
| def get_image_files(self, key: str) -> list[str]: | |
| return self._get_assets(key).image_files | |
| def get_captions(self, key: str) -> list[str]: | |
| return self._get_assets(key).captions | |
| def get_embedding_dim(self, key: str) -> int: | |
| return self._get_assets(key).embedding_dim | |
| def preload(self, key: str) -> None: | |
| """Eagerly load a backbone (useful for warmup in HF Spaces).""" | |
| self._get_assets(key) | |
| def unload_all(self) -> None: | |
| with self._lock: | |
| self._cache.clear() | |
| gc.collect() | |
| # ───── Internals ───── | |
| def _get_assets(self, key: str) -> _BackboneAssets: | |
| if key not in BACKBONES: | |
| raise KeyError( | |
| f"Unknown backbone '{key}'. Known: {list(BACKBONES.keys())}" | |
| ) | |
| with self._lock: | |
| assets = self._cache.get(key) | |
| if assets is not None: | |
| self._cache.move_to_end(key) # touch — most recently used | |
| log.debug("Cache hit: %s", key) | |
| return assets | |
| # Cache miss — load. | |
| log.info("Loading backbone '%s' …", key) | |
| assets = self._load(key) | |
| self._cache[key] = assets | |
| self._cache.move_to_end(key) | |
| # Evict LRU entries beyond the cap. | |
| while len(self._cache) > self._max: | |
| evicted_key, _ = self._cache.popitem(last=False) | |
| log.info("Evicted LRU backbone '%s' (cache size > %d)", | |
| evicted_key, self._max) | |
| gc.collect() | |
| return assets | |
| def _load(self, key: str) -> _BackboneAssets: | |
| spec = BACKBONES[key] | |
| img_h5 = self._h5_dir / f"embeddings_{spec.backbone_id}.h5" | |
| capt_h5 = self._h5_dir / f"captions_{spec.backbone_id}.h5" | |
| img_embeds, img_files = _load_h5(img_h5) | |
| with h5py.File(capt_h5, "r") as f: | |
| capt_embeds = f["embeddings"][:] | |
| captions = [ | |
| s.decode("utf-8") if isinstance(s, bytes) else str(s) | |
| for s in f["captions"][:] | |
| ] | |
| caption_filenames = [ | |
| s.decode("utf-8") if isinstance(s, bytes) else str(s) | |
| for s in f["filenames"][:] | |
| ] | |
| log.info(" %s · %d images · %d captions · dim=%d", | |
| spec.backbone_id, len(img_files), len(captions), | |
| img_embeds.shape[1]) | |
| image_index = _build_faiss_index(img_embeds.copy()) | |
| caption_index = _build_faiss_index(capt_embeds.copy()) | |
| # Local import — keep torch off the import-time hot path. | |
| from embedders import get_embedder | |
| embedder = get_embedder(key) | |
| return _BackboneAssets( | |
| embedder = embedder, | |
| image_index = image_index, | |
| caption_index = caption_index, | |
| image_files = img_files, | |
| captions = captions, | |
| embedding_dim = int(img_embeds.shape[1]), | |
| caption_filenames = caption_filenames, | |
| ) | |
| # ───────────────────────────────────────────────────────────────────── | |
| # Module-level singleton — import this from app.py | |
| # ───────────────────────────────────────────────────────────────────── | |
| manager = BackboneManager() | |
| __all__ = ["manager", "BackboneManager", "BackboneSpec", | |
| "BACKBONES", "DEFAULT_BACKBONE", "MAX_BACKBONES"] | |