#!/usr/bin/env python3 """ Model Loader for Hugging Face Spaces - Robust SAM2 loader with multiple strategies - Correct MatAnyOne loader via official InferenceCore (no transformers) - Clean progress reporting, cleanup, and diagnostics """ from __future__ import annotations import os import gc import time import logging import traceback from pathlib import Path from typing import Optional, Dict, Any, Tuple, Callable import torch from core.exceptions import ModelLoadingError from utils.hardware.device_manager import DeviceManager from utils.system.memory_manager import MemoryManager logger = logging.getLogger(__name__) # ------------------------------ # Data wrapper # ------------------------------ class LoadedModel: def __init__(self, model=None, model_id: str = "", load_time: float = 0.0, device: str = "", framework: str = ""): self.model = model self.model_id = model_id self.load_time = load_time self.device = device self.framework = framework def to_dict(self) -> Dict[str, Any]: return { "model_id": self.model_id, "framework": self.framework, "device": self.device, "load_time": self.load_time, "loaded": self.model is not None, } # ------------------------------ # Loader # ------------------------------ class ModelLoader: def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager): self.device_manager = device_mgr self.memory_manager = memory_mgr self.device = self.device_manager.get_optimal_device() # e.g., cuda:0 or cpu self.sam2_predictor: Optional[LoadedModel] = None self.matanyone_model: Optional[LoadedModel] = None self.checkpoints_dir = "./checkpoints" os.makedirs(self.checkpoints_dir, exist_ok=True) self.loading_stats = { "sam2_load_time": 0.0, "matanyone_load_time": 0.0, "total_load_time": 0.0, "models_loaded": False, "loading_attempts": 0, } logger.info(f"ModelLoader initialized for device: {self.device}") # ---------- Public API ---------- def load_all_models( self, progress_callback: Optional[Callable[[float, str], None]] = None, cancel_event=None ) -> Tuple[Optional[LoadedModel], Optional[LoadedModel]]: """ Loads SAM2 + MatAnyOne. Returns (LoadedModel|None, LoadedModel|None). """ start_time = time.time() self.loading_stats["loading_attempts"] += 1 try: logger.info("Starting model loading process...") if progress_callback: progress_callback(0.0, "Initializing model loading...") self._cleanup_models() # ---- SAM2 ---- logger.info("Loading SAM2 predictor...") if progress_callback: progress_callback(0.1, "Loading SAM2 predictor...") sam2_loaded = self._load_sam2_predictor(progress_callback) if sam2_loaded is None: logger.warning("SAM2 loading failed - a limited fallback will be used at runtime if needed.") else: self.sam2_predictor = sam2_loaded self.loading_stats["sam2_load_time"] = self.sam2_predictor.load_time logger.info(f"SAM2 loaded in {self.loading_stats['sam2_load_time']:.2f}s") # Early exit if cancelled if cancel_event is not None and getattr(cancel_event, "is_set", lambda: False)(): if progress_callback: progress_callback(1.0, "Model loading cancelled") return self.sam2_predictor, None # ---- MatAnyOne ---- logger.info("Loading MatAnyOne model...") if progress_callback: progress_callback(0.6, "Loading MatAnyOne model...") matanyone_loaded = self._load_matanyone(progress_callback) if matanyone_loaded is None: logger.warning("MatAnyOne loading failed - will use simple refinement fallbacks.") else: self.matanyone_model = matanyone_loaded self.loading_stats["matanyone_load_time"] = self.matanyone_model.load_time logger.info(f"MatAnyOne loaded in {self.loading_stats['matanyone_load_time']:.2f}s") # ---- Final status ---- total_time = time.time() - start_time self.loading_stats["total_load_time"] = total_time self.loading_stats["models_loaded"] = bool(self.sam2_predictor or self.matanyone_model) if progress_callback: if self.loading_stats["models_loaded"]: progress_callback(1.0, "Models loaded (fallbacks available if any model failed)") else: progress_callback(1.0, "Using fallback methods (models failed to load)") logger.info(f"Model loading completed in {total_time:.2f}s") return self.sam2_predictor, self.matanyone_model except Exception as e: error_msg = f"Model loading failed: {str(e)}" logger.error(f"{error_msg}\n{traceback.format_exc()}") self._cleanup_models() self.loading_stats["models_loaded"] = False if progress_callback: progress_callback(1.0, f"Error: {error_msg}") return None, None def reload_models(self, progress_callback: Optional[Callable[[float, str], None]] = None) -> Tuple[ Optional[LoadedModel], Optional[LoadedModel] ]: logger.info("Reloading models...") self._cleanup_models() self.loading_stats["models_loaded"] = False return self.load_all_models(progress_callback) @property def models_ready(self) -> bool: return self.sam2_predictor is not None or self.matanyone_model is not None def get_sam2(self): return self.sam2_predictor.model if self.sam2_predictor is not None else None def get_matanyone(self): return self.matanyone_model.model if self.matanyone_model is not None else None def validate_models(self) -> bool: try: ok = False if self.sam2_predictor is not None: model = self.sam2_predictor.model if hasattr(model, "set_image") or hasattr(model, "predict"): ok = True if self.matanyone_model is not None: ok = True return ok except Exception as e: logger.error(f"Model validation failed: {e}") return False def get_model_info(self) -> Dict[str, Any]: info = { "models_loaded": self.loading_stats["models_loaded"], "sam2_loaded": self.sam2_predictor is not None, "matanyone_loaded": self.matanyone_model is not None, "device": str(self.device), "loading_stats": self.loading_stats.copy(), } if self.sam2_predictor is not None: info["sam2_model_type"] = type(self.sam2_predictor.model).__name__ info["sam2_metadata"] = self.sam2_predictor.to_dict() if self.matanyone_model is not None: info["matanyone_model_type"] = type(self.matanyone_model.model).__name__ info["matanyone_metadata"] = self.matanyone_model.to_dict() return info def get_load_summary(self) -> str: if not self.loading_stats["models_loaded"]: return "Models not loaded" sam2_time = self.loading_stats["sam2_load_time"] matanyone_time = self.loading_stats["matanyone_load_time"] total_time = self.loading_stats["total_load_time"] summary = f"Models loaded in {total_time:.1f}s\n" if self.sam2_predictor: summary += f"✓ SAM2: {sam2_time:.1f}s (ID: {self.sam2_predictor.model_id})\n" else: summary += "✗ SAM2: Failed (using fallback)\n" if self.matanyone_model: summary += f"✓ MatAnyOne: {matanyone_time:.1f}s (ID: {self.matanyone_model.model_id})\n" else: summary += "✗ MatAnyOne: Failed (using simple refinement)\n" summary += f"Device: {self.device}" return summary def cleanup(self): self._cleanup_models() logger.info("ModelLoader cleanup completed") # ---------- Internal: SAM2 ---------- def _load_sam2_predictor(self, progress_callback: Optional[Callable[[float, str], None]] = None) -> Optional[LoadedModel]: """ Try multiple SAM2 loading strategies: official -> transformers -> dummy fallback. """ # Choose model size heuristically model_size = "large" try: if hasattr(self.device_manager, "get_device_memory_gb"): memory_gb = self.device_manager.get_device_memory_gb() if memory_gb < 4: model_size = "tiny" elif memory_gb < 8: model_size = "small" elif memory_gb < 12: model_size = "base" logger.info(f"Selected SAM2 {model_size} based on {memory_gb}GB VRAM") except Exception as e: logger.warning(f"Could not determine device memory: {e}") model_size = "tiny" model_map = { "tiny": "facebook/sam2.1-hiera-tiny", "small": "facebook/sam2.1-hiera-small", "base": "facebook/sam2.1-hiera-base-plus", "large": "facebook/sam2.1-hiera-large", } model_id = model_map.get(model_size, model_map["tiny"]) if progress_callback: progress_callback(0.3, f"Loading SAM2 ({model_size})...") methods = [ ("official", self._try_load_sam2_official, model_id), ("direct", self._try_load_sam2_direct, model_id), ("manual", self._try_load_sam2_manual, model_id), ] for name, fn, mid in methods: try: logger.info(f"Attempting SAM2 load via {name} method ({mid})...") result = fn(mid) if result is not None: logger.info(f"SAM2 loaded successfully via {name} method") return result except Exception as e: logger.error(f"SAM2 {name} method failed: {e}") logger.debug(traceback.format_exc()) continue logger.error("All SAM2 loading methods failed") return None def _try_load_sam2_official(self, model_id: str) -> Optional[LoadedModel]: """ Official predictor path (Meta's SAM2ImagePredictor). """ from sam2.sam2_image_predictor import SAM2ImagePredictor # Space-specific hub flags os.environ["HF_HUB_DISABLE_SYMLINKS"] = "1" os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" cache_dir = os.path.join(self.checkpoints_dir, "sam2_cache") os.makedirs(cache_dir, exist_ok=True) t0 = time.time() predictor = SAM2ImagePredictor.from_pretrained( model_id, cache_dir=cache_dir, local_files_only=False, trust_remote_code=True, ) if hasattr(predictor, "model"): predictor.model = predictor.model.to(self.device) t1 = time.time() return LoadedModel( model=predictor, model_id=model_id, load_time=t1 - t0, device=str(self.device), framework="sam2" ) def _try_load_sam2_direct(self, model_id: str) -> Optional[LoadedModel]: """ Transformers AutoModel path (best-effort; API may vary). """ from transformers import AutoModel, AutoProcessor t0 = time.time() model = AutoModel.from_pretrained( model_id, trust_remote_code=True, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, ).to(self.device) try: processor = AutoProcessor.from_pretrained(model_id) except Exception: processor = None t1 = time.time() class SAM2Wrapper: def __init__(self, model, processor=None): self.model = model self.processor = processor def set_image(self, image): self.current_image = image def predict(self, *args, **kwargs): return self.model(*args, **kwargs) wrapped = SAM2Wrapper(model, processor) return LoadedModel( model=wrapped, model_id=model_id, load_time=t1 - t0, device=str(self.device), framework="sam2-transformers", ) def _try_load_sam2_manual(self, model_id: str) -> Optional[LoadedModel]: """ Dummy fallback that won't crash the app. """ class DummySAM2: def __init__(self, device): self.device = device self.model = None def set_image(self, image): self.current_image = image def predict(self, point_coords=None, point_labels=None, box=None, **kwargs): import numpy as np if hasattr(self, "current_image"): h, w = self.current_image.shape[:2] else: h, w = 512, 512 return { "masks": np.ones((1, h, w), dtype=np.float32), "scores": np.array([0.5]), "logits": np.ones((1, h, w), dtype=np.float32), } t0 = time.time() dummy = DummySAM2(self.device) t1 = time.time() logger.warning("Using manual SAM2 fallback (limited functionality)") return LoadedModel( model=dummy, model_id=f"{model_id}-fallback", load_time=t1 - t0, device=str(self.device), framework="sam2-fallback" ) # ---------- Internal: MatAnyOne ---------- def _load_matanyone(self, progress_callback: Optional[Callable[[float, str], None]] = None) -> Optional[LoadedModel]: """ Correct MatAnyOne loader using official package API. """ if progress_callback: progress_callback(0.7, "Loading MatAnyOne (InferenceCore)...") try: return self._try_load_matanyone_official() except Exception as e: logger.error(f"MatAnyOne official loader failed: {e}") logger.debug(traceback.format_exc()) logger.warning("Falling back to simple MatAnyOne placeholder.") return self._try_load_matanyone_fallback() def _try_load_matanyone_official(self) -> Optional[LoadedModel]: """ Official MatAnyOne via package's InferenceCore. IMPORTANT: pass model id POSITIONALLY; do NOT use repo_id= or transformers. """ from matanyone import InferenceCore t0 = time.time() processor = InferenceCore("PeiqingYang/MatAnyone") t1 = time.time() return LoadedModel( model=processor, model_id="PeiqingYang/MatAnyone", load_time=t1 - t0, device=str(self.device), framework="matanyone", ) def _try_load_matanyone_fallback(self) -> Optional[LoadedModel]: """ Minimal placeholder that safely passes masks through. """ class FallbackMatAnyone: def __init__(self, device): self.device = device def process(self, image, mask): # Identity pass-through (keeps pipeline alive) return mask t0 = time.time() model = FallbackMatAnyone(self.device) t1 = time.time() logger.warning("Using MatAnyOne fallback (limited functionality)") return LoadedModel( model=model, model_id="MatAnyone-fallback", load_time=t1 - t0, device=str(self.device), framework="matanyone-fallback" ) # ---------- Internal: cleanup ---------- def _cleanup_models(self): if self.sam2_predictor is not None: del self.sam2_predictor self.sam2_predictor = None if self.matanyone_model is not None: del self.matanyone_model self.matanyone_model = None if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() logger.debug("Model cleanup completed")