142 lines
5.0 KiB
Python
142 lines
5.0 KiB
Python
"""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
|