| """ |
| 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] |
|
|
| |
| if model_info["loader"] == "inference_api": |
| return True |
|
|
| |
| if model_info["status"] == "placeholder": |
| logger.warning(f"Model {name} is a placeholder — not yet available") |
| return False |
|
|
| |
| if model_info["status"] == "loaded" and model_info["instance"] is not None: |
| return True |
|
|
| |
| if self._active_local_model and self._active_local_model != name: |
| self._unload_active_local() |
|
|
| |
| 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") |
|
|
|
|
| |
| registry = ModelRegistry() |
|
|
