first version

2026-04-25 16:23:09 +05:30
commit 7d3c36050b
73 changed files with 19033 additions and 0 deletions
--- a/sam2-backend/registry.py
+++ b/sam2-backend/registry.py
@@ -0,0 +1,141 @@
+"""Model registry. Reads models.yaml, lazy-loads adapters, caches instances.
+
+Only one adapter is resident at a time by default — switching models unloads
+the previous one to keep GPU memory bounded. Set SAM_BACKEND_KEEP_ALL=1 to
+disable eviction (useful if you have VRAM headroom and switch often).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import yaml
+
+from adapters import ADAPTERS, BaseAdapter, PredictResult
+
+log = logging.getLogger("sam2-backend.registry")
+
+
+@dataclass
+class ModelSpec:
+    id: str
+    name: str
+    family: str
+    config: str
+    checkpoint: str
+    # If true, `checkpoint` is treated as an opaque identifier (e.g. a
+    # HuggingFace model id like "facebook/sam3") rather than a file on disk.
+    # Availability check is skipped for remote entries.
+    remote: bool = False
+    available: bool = False
+    # Resolved absolute paths for logging / responses.
+    checkpoint_abs: str = ""
+
+
+@dataclass
+class Registry:
+    specs: dict[str, ModelSpec] = field(default_factory=dict)
+    default_id: Optional[str] = None
+    device: str = "cpu"
+    keep_all: bool = False
+    _loaded: dict[str, BaseAdapter] = field(default_factory=dict)
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def list_specs(self) -> list[ModelSpec]:
+        return list(self.specs.values())
+
+    def resolve(self, model_id: Optional[str]) -> ModelSpec:
+        if not model_id:
+            if self.default_id is None:
+                raise KeyError("no default model configured")
+            model_id = self.default_id
+        if model_id not in self.specs:
+            raise KeyError(f"unknown model: {model_id}")
+        return self.specs[model_id]
+
+    def get_adapter(self, model_id: str) -> BaseAdapter:
+        """Returns a loaded adapter for `model_id`, loading on first use."""
+        with self._lock:
+            if model_id in self._loaded:
+                return self._loaded[model_id]
+            spec = self.resolve(model_id)
+            if not spec.available:
+                raise FileNotFoundError(
+                    f"checkpoint missing: {spec.checkpoint_abs}"
+                )
+            cls = ADAPTERS.get(spec.family)
+            if cls is None:
+                raise ValueError(
+                    f"no adapter registered for family '{spec.family}'"
+                )
+            if not self.keep_all:
+                # Evict whichever model is currently resident.
+                for old_id, adapter in list(self._loaded.items()):
+                    log.info("evicting %s to make room for %s", old_id, model_id)
+                    adapter.unload()
+                    del self._loaded[old_id]
+            adapter = cls()
+            adapter.load(spec.config, spec.checkpoint_abs, self.device)
+            self._loaded[model_id] = adapter
+            return adapter
+
+    def predict_bbox(self, model_id: Optional[str], image_rgb, bbox_xyxy) -> tuple[ModelSpec, PredictResult]:
+        spec = self.resolve(model_id)
+        adapter = self.get_adapter(spec.id)
+        # Serialize inference: most of these adapters aren't reentrant.
+        with self._lock:
+            result = adapter.predict_bbox(image_rgb, bbox_xyxy)
+        return spec, result
+
+
+def load_registry(models_yaml: Path, device: str) -> Registry:
+    with open(models_yaml, "r") as f:
+        data = yaml.safe_load(f)
+
+    base_dir = models_yaml.parent
+    keep_all = os.environ.get("SAM_BACKEND_KEEP_ALL", "0") == "1"
+    reg = Registry(device=device, keep_all=keep_all)
+    reg.default_id = data.get("default")
+
+    for entry in data.get("models", []):
+        remote = bool(entry.get("remote", False))
+        ckpt_str = entry["checkpoint"]
+        if remote:
+            # Opaque identifier — leave as-is, mark available so the UI
+            # shows it. The adapter resolves/downloads on first use.
+            ckpt_abs = ckpt_str
+            available = True
+        else:
+            ckpt = Path(ckpt_str)
+            ckpt_path = (ckpt if ckpt.is_absolute() else (base_dir / ckpt)).resolve()
+            ckpt_abs = str(ckpt_path)
+            available = ckpt_path.is_file()
+        spec = ModelSpec(
+            id=entry["id"],
+            name=entry["name"],
+            family=entry["family"],
+            config=entry["config"],
+            checkpoint=ckpt_str,
+            remote=remote,
+            available=available,
+            checkpoint_abs=ckpt_abs,
+        )
+        reg.specs[spec.id] = spec
+        log.info(
+            "registered model %s (available=%s, ckpt=%s)",
+            spec.id, spec.available, spec.checkpoint_abs,
+        )
+
+    # Fall back the default to the first *available* model if the configured
+    # default's checkpoint is missing.
+    if reg.default_id and not reg.specs.get(reg.default_id, ModelSpec("", "", "", "", "")).available:
+        for s in reg.specs.values():
+            if s.available:
+                reg.default_id = s.id
+                break
+    return reg