Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -1,12 +1,15 @@
 #!/usr/bin/env python3
 """
-FIXED Model Loading Module for HuggingFace Spaces
-Handles the list index out of range error
+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 sys
 import time
 import logging
 import traceback
@@ -21,6 +24,10 @@
 
 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
@@ -29,40 +36,51 @@ def __init__(self, model=None, model_id: str = "", load_time: float = 0.0, devic
         self.device = device
         self.framework = framework
 
-    def to_dict(self):
+    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
+            "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()
+        self.device = self.device_manager.get_optimal_device()  # e.g., cuda:0 or cpu
 
-        self.sam2_predictor = None
-        self.matanyone_model = None
+        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
+            "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}")
 
-    def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_event=None) -> Tuple[Any, Any]:
+    # ---------- 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
+        self.loading_stats["loading_attempts"] += 1
 
         try:
             logger.info("Starting model loading process...")
@@ -71,66 +89,141 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
 
             self._cleanup_models()
 
-            # Load SAM2 with better error handling
+            # ---- SAM2 ----
             logger.info("Loading SAM2 predictor...")
             if progress_callback:
                 progress_callback(0.1, "Loading SAM2 predictor...")
-            sam2_loaded = self._load_sam2_predictor_safe(progress_callback)
+            sam2_loaded = self._load_sam2_predictor(progress_callback)
 
             if sam2_loaded is None:
-                logger.warning("SAM2 loading failed - will use fallback segmentation")
+                logger.warning("SAM2 loading failed - a limited fallback will be used at runtime if needed.")
             else:
                 self.sam2_predictor = sam2_loaded
-                sam2_time = self.sam2_predictor.load_time
-                self.loading_stats['sam2_load_time'] = sam2_time
-                logger.info(f"SAM2 loaded in {sam2_time:.2f}s")
+                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
 
-            # Load MatAnyOne with better error handling
+            # ---- MatAnyOne ----
             logger.info("Loading MatAnyOne model...")
             if progress_callback:
                 progress_callback(0.6, "Loading MatAnyOne model...")
-            
-            matanyone_loaded = self._load_matanyone_model_safe(progress_callback)
+            matanyone_loaded = self._load_matanyone(progress_callback)
 
             if matanyone_loaded is None:
-                logger.warning("MatAnyOne loading failed - will use OpenCV refinement")
+                logger.warning("MatAnyOne loading failed - will use simple refinement fallbacks.")
             else:
                 self.matanyone_model = matanyone_loaded
-                matanyone_time = self.matanyone_model.load_time
-                self.loading_stats['matanyone_load_time'] = matanyone_time
-                logger.info(f"MatAnyOne loaded in {matanyone_time:.1f}s")
+                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
+            # ---- 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)
+            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.sam2_predictor or self.matanyone_model:
-                    progress_callback(1.0, "Models loaded (with fallbacks available)")
+                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)
+            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
+            self.loading_stats["models_loaded"] = False
             if progress_callback:
                 progress_callback(1.0, f"Error: {error_msg}")
             return None, None
 
-    def _load_sam2_predictor_safe(self, progress_callback: Optional[Callable] = None):
-        """Load SAM2 with comprehensive error handling for HuggingFace Spaces"""
-        
-        # Determine model size based on available memory
+    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'):
+            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"
@@ -138,290 +231,203 @@ def _load_sam2_predictor_safe(self, progress_callback: Optional[Callable] = None
                     model_size = "small"
                 elif memory_gb < 12:
                     model_size = "base"
-                logger.info(f"Selected SAM2 {model_size} based on {memory_gb}GB memory")
+                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"  # Default to tiny for Spaces
+            model_size = "tiny"
 
         model_map = {
             "tiny": "facebook/sam2.1-hiera-tiny",
-            "small": "facebook/sam2.1-hiera-small", 
+            "small": "facebook/sam2.1-hiera-small",
             "base": "facebook/sam2.1-hiera-base-plus",
-            "large": "facebook/sam2.1-hiera-large"
+            "large": "facebook/sam2.1-hiera-large",
         }
         model_id = model_map.get(model_size, model_map["tiny"])
-        
-        logger.info(f"[DIAG] Loading SAM2 model_id: {model_id} on device {self.device}")
 
         if progress_callback:
-            progress_callback(0.3, f"Loading SAM2 {model_size} model...")
+            progress_callback(0.3, f"Loading SAM2 ({model_size})...")
 
-        # Try multiple loading strategies
-        loading_methods = [
+        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 method_name, method_func, model_id in loading_methods:
+        for name, fn, mid in methods:
             try:
-                logger.info(f"Attempting SAM2 load via {method_name} method...")
-                result = method_func(model_id)
+                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 {method_name} method")
+                    logger.info(f"SAM2 loaded successfully via {name} method")
                     return result
-            except IndexError as e:
-                logger.error(f"SAM2 {method_name} method - IndexError: {e}")
-                logger.debug(f"Full traceback:\n{traceback.format_exc()}")
-                continue
             except Exception as e:
-                logger.error(f"SAM2 {method_name} method failed: {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):
-        """Try the official from_pretrained method"""
-        try:
-            from sam2.sam2_image_predictor import SAM2ImagePredictor
-            
-            # Set environment variables that might help in Spaces
-            os.environ['HF_HUB_DISABLE_SYMLINKS'] = '1'
-            os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '0'
-            
-            t0 = time.time()
-            
-            # Try with explicit cache directory
-            cache_dir = os.path.join(self.checkpoints_dir, "sam2_cache")
-            os.makedirs(cache_dir, exist_ok=True)
-            
-            # Log what we're about to do
-            logger.debug(f"Calling SAM2ImagePredictor.from_pretrained('{model_id}')")
-            
-            # This is where the IndexError likely happens
-            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"
-            )
-        except Exception as e:
-            logger.error(f"Official SAM2 loading failed: {e}")
-            raise
-
-    def _try_load_sam2_direct(self, model_id: str):
-        """Try loading SAM2 using transformers AutoModel"""
-        try:
-            from transformers import AutoModel, AutoProcessor
-            
-            t0 = time.time()
-            
-            # Try loading as a standard transformers model
-            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 to get processor
-            try:
-                processor = AutoProcessor.from_pretrained(model_id)
-            except:
-                processor = None
-                
-            t1 = time.time()
-            
-            # Wrap in a compatible interface
-            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):
-                    # Basic prediction interface
-                    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"
-            )
-        except Exception as e:
-            logger.error(f"Direct SAM2 loading failed: {e}")
-            raise
-
-    def _try_load_sam2_manual(self, model_id: str):
-        """Try manual model construction as last resort"""
-        try:
-            # This is a fallback - create a dummy model that at least won't crash
-            logger.warning("Using manual SAM2 construction (limited functionality)")
-            
-            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):
-                    # Return a dummy mask
-                    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)
-                    }
-            
-            dummy = DummySAM2(self.device)
-            
-            return LoadedModel(
-                model=dummy,
-                model_id=f"{model_id}-fallback",
-                load_time=0.1,
-                device=str(self.device),
-                framework="sam2-fallback"
-            )
-        except Exception as e:
-            logger.error(f"Manual SAM2 construction failed: {e}")
-            raise
-
-    def _load_matanyone_model_safe(self, progress_callback: Optional[Callable] = None):
-        """Load MatAnyOne with comprehensive error handling"""
-        
-        loading_methods = [
-            ("official", self._try_load_matanyone_official),
-            ("alternative", self._try_load_matanyone_alternative),
-            ("fallback", self._try_load_matanyone_fallback),
-        ]
-
-        for method_name, method_func in loading_methods:
-            try:
-                logger.info(f"Attempting MatAnyOne load via {method_name} method...")
-                result = method_func(progress_callback)
-                if result is not None:
-                    logger.info(f"MatAnyOne loaded successfully via {method_name} method")
-                    return result
-            except IndexError as e:
-                logger.error(f"MatAnyOne {method_name} method - IndexError: {e}")
-                logger.debug(f"Full traceback:\n{traceback.format_exc()}")
-                continue
-            except Exception as e:
-                logger.error(f"MatAnyOne {method_name} method failed: {e}")
-                continue
+    def _try_load_sam2_official(self, model_id: str) -> Optional[LoadedModel]:
+        """
+        Official predictor path (Meta's SAM2ImagePredictor).
+        """
+        from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-        logger.error("All MatAnyOne loading methods failed")
-        return None
+        # Space-specific hub flags
+        os.environ["HF_HUB_DISABLE_SYMLINKS"] = "1"
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
-    def _try_load_matanyone_official(self, progress_callback):
-        """Try the official MatAnyOne loading method"""
-        if progress_callback:
-            progress_callback(0.7, "Loading MatAnyOne model (official)...")
+        cache_dir = os.path.join(self.checkpoints_dir, "sam2_cache")
+        os.makedirs(cache_dir, exist_ok=True)
 
-        from matanyone import InferenceCore
         t0 = time.time()
-        
-        # Set cache directory
-        cache_dir = os.path.join(self.checkpoints_dir, "matanyone_cache")
-        os.makedirs(cache_dir, exist_ok=True)
-        
-        processor = InferenceCore(
-            repo_id="PeiqingYang/MatAnyone",
-            device=self.device,
-            dtype=torch.float32,
-            cache_dir=cache_dir
+        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=processor,
-            model_id="PeiqingYang/MatAnyone",
-            load_time=t1-t0,
-            device=str(self.device),
-            framework="matanyone"
+            model=predictor, model_id=model_id, load_time=t1 - t0, device=str(self.device), framework="sam2"
         )
 
-    def _try_load_matanyone_alternative(self, progress_callback):
-        """Try alternative loading for MatAnyOne"""
-        if progress_callback:
-            progress_callback(0.7, "Loading MatAnyOne model (alternative)...")
-            
-        # Try loading via transformers
-        from transformers import AutoModel
-        
+    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(
-            "PeiqingYang/MatAnyone",
+            model_id,
             trust_remote_code=True,
-            torch_dtype=torch.float32
+            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()
-        
-        # Wrap for compatibility
-        class MatAnyoneWrapper:
-            def __init__(self, model):
+
+        class SAM2Wrapper:
+            def __init__(self, model, processor=None):
                 self.model = model
-                
-            def process(self, image, mask):
-                return self.model(image, mask)
-        
+                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=MatAnyoneWrapper(model),
-            model_id="PeiqingYang/MatAnyone-alt",
-            load_time=t1-t0,
+            model=wrapped,
+            model_id=model_id,
+            load_time=t1 - t0,
             device=str(self.device),
-            framework="matanyone-transformers"
+            framework="sam2-transformers",
         )
 
-    def _try_load_matanyone_fallback(self, progress_callback):
-        """Create a fallback MatAnyOne that won't crash"""
+    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, "Using MatAnyOne fallback...")
-            
-        logger.warning("Using fallback MatAnyOne (limited functionality)")
-        
+            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):
-                # Just return the mask unchanged
+                # 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=FallbackMatAnyone(self.device),
-            model_id="MatAnyone-fallback",
-            load_time=0.1,
-            device=str(self.device),
-            framework="matanyone-fallback"
+            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
@@ -433,77 +439,3 @@ def _cleanup_models(self):
             torch.cuda.empty_cache()
         gc.collect()
         logger.debug("Model cleanup completed")
-
-    def cleanup(self):
-        self._cleanup_models()
-        logger.info("ModelLoader cleanup completed")
-
-    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 += f"✗ 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 += f"✗ MatAnyOne: Failed (using OpenCV)\n"
-        summary += f"Device: {self.device}"
-        return summary
-
-    def get_matanyone(self):
-        # Return the actual model from inside the LoadedModel wrapper
-        if self.matanyone_model is not None:
-            return self.matanyone_model.model if hasattr(self.matanyone_model, 'model') else None
-        return None
-
-    def get_sam2(self):
-        # Return the actual model from inside the LoadedModel wrapper
-        if self.sam2_predictor is not None:
-            return self.sam2_predictor.model if hasattr(self.sam2_predictor, 'model') else None
-        return None
-
-    def validate_models(self) -> bool:
-        try:
-            has_valid_model = False
-            if self.sam2_predictor is not None:
-                model = self.sam2_predictor.model
-                if hasattr(model, 'set_image') or hasattr(model, 'predict'):
-                    has_valid_model = True
-            if self.matanyone_model is not None:
-                has_valid_model = True
-            return has_valid_model
-        except Exception as e:
-            logger.error(f"Model validation failed: {e}")
-            return False
-
-    def reload_models(self, progress_callback: Optional[Callable] = None) -> Tuple[Any, Any]:
-        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