""" Model Registry for NeuroSAM3 — Multi-model management with lazy loading. Manages loading/unloading of multiple models with GPU memory awareness. Only one large model is loaded at a time to avoid OOM on HF Spaces. """ from typing import Optional, Dict, Any import torch from logger_config import logger from config import ( SAM_MODEL_ID, MEDSAM_MODEL_ID, BIOMEDCLIP_MODEL_ID, UNET_TUMOR_MODEL_ID, LLM_MODELS, HF_TOKEN, ) class ModelRegistry: """Manages loading/unloading of multiple models with GPU memory awareness.""" def __init__(self): self._models: Dict[str, Dict[str, Any]] = { "sam3": { "model_id": SAM_MODEL_ID, "loader": "local", "instance": None, "processor": None, "status": "unloaded", "category": "segmentation", }, "medsam": { "model_id": MEDSAM_MODEL_ID, "loader": "local", "instance": None, "processor": None, "status": "unloaded", "category": "segmentation", }, "biomedclip": { "model_id": BIOMEDCLIP_MODEL_ID, "loader": "local", "instance": None, "processor": None, "status": "unloaded", "category": "classification", }, "unet_tumor": { "model_id": UNET_TUMOR_MODEL_ID, "loader": "local", "instance": None, "processor": None, "status": "placeholder", "category": "segmentation", }, "monai_brats": { "model_id": "monai/brats_mri_segmentation", "loader": "local", "instance": None, "processor": None, "status": "unloaded", "category": "segmentation", }, "gemma": { "model_id": LLM_MODELS["gemma"], "loader": "inference_api", "instance": None, "processor": None, "status": "ready", "category": "reasoning", }, "kimi": { "model_id": LLM_MODELS["kimi"], "loader": "inference_api", "instance": None, "processor": None, "status": "ready", "category": "reasoning", }, "glm": { "model_id": LLM_MODELS["glm"], "loader": "inference_api", "instance": None, "processor": None, "status": "ready", "category": "reasoning", }, } self._active_local_model: Optional[str] = None def get_status(self, name: str) -> str: """Get model status.""" if name not in self._models: return "unknown" return self._models[name]["status"] def is_loaded(self, name: str) -> bool: """Check if a model is loaded and ready.""" return self._models.get(name, {}).get("status") == "loaded" def is_available(self, name: str) -> bool: """Check if a model is available (loaded or API-ready).""" status = self.get_status(name) return status in ("loaded", "ready") def get_model(self, name: str) -> Optional[Any]: """Get model instance (None if not loaded).""" return self._models.get(name, {}).get("instance") def get_processor(self, name: str) -> Optional[Any]: """Get model processor (None if not loaded).""" return self._models.get(name, {}).get("processor") def get_model_id(self, name: str) -> Optional[str]: """Get HuggingFace model ID.""" return self._models.get(name, {}).get("model_id") def list_models(self) -> Dict[str, Dict[str, str]]: """List all models with their status and category.""" return { name: { "model_id": info["model_id"], "status": info["status"], "category": info["category"], "loader": info["loader"], } for name, info in self._models.items() } def list_available(self, category: Optional[str] = None) -> list: """List available model names, optionally filtered by category.""" models = [] for name, info in self._models.items(): if info["status"] in ("loaded", "ready", "unloaded"): if category is None or info["category"] == category: models.append(name) return models def _unload_active_local(self): """Unload the currently active local model to free GPU memory.""" if self._active_local_model is None: return name = self._active_local_model model_info = self._models[name] if model_info["instance"] is not None: logger.info(f"Unloading {name} to free GPU memory...") try: model_info["instance"].to("cpu") del model_info["instance"] model_info["instance"] = None model_info["processor"] = None model_info["status"] = "unloaded" torch.cuda.empty_cache() logger.info(f"Unloaded {name}") except Exception as e: logger.warning(f"Error unloading {name}: {e}") self._active_local_model = None def load_model(self, name: str) -> bool: """ Load a model to GPU. Unloads other local models if needed. Returns True if model is ready to use, False otherwise. """ if name not in self._models: logger.error(f"Unknown model: {name}") return False model_info = self._models[name] # API models are always "ready" if model_info["loader"] == "inference_api": return True # Placeholder models can't be loaded if model_info["status"] == "placeholder": logger.warning(f"Model {name} is a placeholder — not yet available") return False # Already loaded if model_info["status"] == "loaded" and model_info["instance"] is not None: return True # Unload current model if different if self._active_local_model and self._active_local_model != name: self._unload_active_local() # Load the requested model try: logger.info(f"Loading model: {name} ({model_info['model_id']})") if name == "sam3": success = self._load_sam3() elif name == "medsam": success = self._load_medsam() elif name == "biomedclip": success = self._load_biomedclip() elif name == "monai_brats": success = self._load_monai_brats() else: logger.error(f"No loader implemented for: {name}") return False if success: model_info["status"] = "loaded" self._active_local_model = name logger.info(f"Model {name} loaded successfully") return success except Exception as e: logger.error(f"Failed to load {name}: {e}", exc_info=True) model_info["status"] = "error" return False def _load_sam3(self) -> bool: """Load SAM3 model (delegated to existing models.py logic).""" from models import initialize_model, get_model, get_processor success = initialize_model() if success: self._models["sam3"]["instance"] = get_model() self._models["sam3"]["processor"] = get_processor() return success def _load_medsam(self) -> bool: """Load MedSAM model.""" from medsam_inference import load_medsam_model model, processor = load_medsam_model() if model is not None: self._models["medsam"]["instance"] = model self._models["medsam"]["processor"] = processor return True return False def _load_biomedclip(self) -> bool: """Load BiomedCLIP model.""" from biomedclip_inference import load_biomedclip_model model, preprocess, tokenizer = load_biomedclip_model() if model is not None: self._models["biomedclip"]["instance"] = model self._models["biomedclip"]["processor"] = { "preprocess": preprocess, "tokenizer": tokenizer, } return True return False def _load_monai_brats(self) -> bool: """Load MONAI BraTS SegResNet model.""" from monai_inference import load_brats_model, is_brats_model_loaded success = load_brats_model() if success and is_brats_model_loaded(): from monai_inference import _brats_model self._models["monai_brats"]["instance"] = _brats_model return True return False def register_sam3_from_existing(self, model, processor): """Register an already-loaded SAM3 model (from app.py startup).""" if model is not None and processor is not None: self._models["sam3"]["instance"] = model self._models["sam3"]["processor"] = processor self._models["sam3"]["status"] = "loaded" self._active_local_model = "sam3" logger.info("Registered existing SAM3 model in registry") # Global singleton registry = ModelRegistry()