Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -1,11 +1,7 @@
 #!/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
-- NEW: Global MatAnyOne step/process shape guard to prevent 5D tensors
-- UPDATED: Enhanced MatAnyone wrapper support for component masks
+Unified Model Loader
+Coordinates separate SAM2 and MatAnyone loaders for cleaner architecture
 """
 
 from __future__ import annotations
@@ -14,8 +10,6 @@
 import gc
 import time
 import logging
-import traceback
-from pathlib import Path
 from typing import Optional, Dict, Any, Tuple, Callable
 
 import torch
@@ -24,14 +18,17 @@
 from utils.hardware.device_manager import DeviceManager
 from utils.system.memory_manager import MemoryManager
 
+# Import the specialized loaders
+from models.loaders.sam2_loader import SAM2Loader
+from models.loaders.matanyone_loader import MatAnyoneLoader
+
 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 = ""):
+    """Container for loaded model information"""
+    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
@@ -48,22 +45,23 @@ def to_dict(self) -> Dict[str, Any]:
         }
 
 
-# ------------------------------
-# Loader
-# ------------------------------
 class ModelLoader:
+    """Main model loader that coordinates SAM2 and MatAnyone loaders"""
+    
     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.device = self.device_manager.get_optimal_device()
+        
+        # Initialize specialized loaders
+        self.sam2_loader = SAM2Loader(device=str(self.device))
+        self.matanyone_loader = MatAnyoneLoader(device=str(self.device))
+        
+        # Model storage
         self.sam2_predictor: Optional[LoadedModel] = None
         self.matanyone_model: Optional[LoadedModel] = None
-        self._matanyone_wrapper = None  # Cache for enhanced wrapper
-
-        self.checkpoints_dir = "./checkpoints"
-        os.makedirs(self.checkpoints_dir, exist_ok=True)
-
+        
+        # Statistics
         self.loading_stats = {
             "sam2_load_time": 0.0,
             "matanyone_load_time": 0.0,
@@ -71,85 +69,114 @@ def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
             "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
+        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).
+        Load all models using specialized loaders
+        
+        Args:
+            progress_callback: Optional callback for progress updates
+            cancel_event: Optional threading.Event for cancellation
+            
+        Returns:
+            Tuple of (sam2_model, matanyone_model)
         """
         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...")
-
+                
+            # Clean up any existing models
             self._cleanup_models()
-
-            # ---- SAM2 ----
-            logger.info("Loading SAM2 predictor...")
+            
+            # Load SAM2
             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.")
+                progress_callback(0.1, "Loading SAM2 model...")
+            
+            sam2_start = time.time()
+            sam2_model = self.sam2_loader.load()
+            sam2_time = time.time() - sam2_start
+            
+            if sam2_model:
+                self.sam2_predictor = LoadedModel(
+                    model=sam2_model,
+                    model_id=self.sam2_loader.model_id,
+                    load_time=sam2_time,
+                    device=str(self.device),
+                    framework="sam2"
+                )
+                self.loading_stats["sam2_load_time"] = sam2_time
+                logger.info(f"SAM2 loaded in {sam2_time:.2f}s")
             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)():
+                logger.warning("SAM2 loading failed")
+                
+            # Check for cancellation
+            if cancel_event and cancel_event.is_set():
                 if progress_callback:
                     progress_callback(1.0, "Model loading cancelled")
                 return self.sam2_predictor, None
-
-            # ---- MatAnyOne ----
-            logger.info("Loading MatAnyOne model...")
+                
+            # Load MatAnyone
             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.")
+                progress_callback(0.6, "Loading MatAnyone model...")
+                
+            matanyone_start = time.time()
+            matanyone_model = self.matanyone_loader.load()
+            matanyone_time = time.time() - matanyone_start
+            
+            if matanyone_model:
+                self.matanyone_model = LoadedModel(
+                    model=matanyone_model,
+                    model_id=self.matanyone_loader.model_id,
+                    load_time=matanyone_time,
+                    device=str(self.device),
+                    framework="matanyone"
+                )
+                self.loading_stats["matanyone_load_time"] = matanyone_time
+                logger.info(f"MatAnyone loaded in {matanyone_time:.2f}s")
             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 ----
+                logger.warning("MatAnyone loading failed")
+                
+            # Update statistics
             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)
-
+            
+            # Final progress update
             if progress_callback:
                 if self.loading_stats["models_loaded"]:
-                    progress_callback(1.0, "Models loaded (fallbacks available if any model failed)")
+                    progress_callback(1.0, "Models loaded successfully")
                 else:
-                    progress_callback(1.0, "Using fallback methods (models failed to load)")
-
+                    progress_callback(1.0, "Model loading completed with failures")
+                    
             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()}")
+            logger.error(error_msg)
             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]
-    ]:
+    def reload_models(
+        self,
+        progress_callback: Optional[Callable[[float, str], None]] = None
+    ) -> Tuple[Optional[LoadedModel], Optional[LoadedModel]]:
+        """Reload all models from scratch"""
         logger.info("Reloading models...")
         self._cleanup_models()
         self.loading_stats["models_loaded"] = False
@@ -157,486 +184,106 @@ def reload_models(self, progress_callback: Optional[Callable[[float, str], None]
 
     @property
     def models_ready(self) -> bool:
+        """Check if any models are loaded and ready"""
         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
+        """Get SAM2 predictor model"""
+        return self.sam2_predictor.model if self.sam2_predictor else None
 
     def get_matanyone(self):
-        """Get MatAnyone processor, optionally wrapped with enhanced features."""
-        if self.matanyone_model is None:
-            return None
-        
-        # Check if we should use the enhanced wrapper
-        try:
-            from app_config import get_config
-            config = get_config()
-            
-            if config.matanyone_enabled and (config.use_component_masks or 
-                                            config.matanyone_edge_enhancement or 
-                                            config.matanyone_hair_refinement):
-                # Use enhanced wrapper for advanced features
-                try:
-                    from models.wrappers.matanyone_wrapper import MatAnyOneWrapper
-                    
-                    if self._matanyone_wrapper is None:
-                        self._matanyone_wrapper = MatAnyOneWrapper(
-                            self.matanyone_model.model,
-                            device=self.device,
-                            config=config.get_matanyone_config()
-                        )
-                        logger.info("Using enhanced MatAnyone wrapper with component support")
-                    return self._matanyone_wrapper
-                except ImportError as e:
-                    logger.warning(f"Enhanced MatAnyone wrapper not available: {e}")
-                except Exception as e:
-                    logger.error(f"Failed to initialize enhanced MatAnyone wrapper: {e}")
-                    
-        except Exception as e:
-            logger.debug(f"Could not check for enhanced wrapper configuration: {e}")
-        
-        # Return raw model for basic usage
-        return self.matanyone_model.model if self.matanyone_model is not None else None
+        """Get MatAnyone processor model"""
+        return self.matanyone_model.model if self.matanyone_model else None
 
     def validate_models(self) -> bool:
+        """Validate that loaded models have expected interfaces"""
         try:
-            ok = False
-            if self.sam2_predictor is not None:
+            valid = False
+            
+            if self.sam2_predictor:
                 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
+                if hasattr(model, "set_image") and hasattr(model, "predict"):
+                    valid = True
+                    logger.info("SAM2 model validated")
+                    
+            if self.matanyone_model:
+                model = self.matanyone_model.model
+                if hasattr(model, "step") or hasattr(model, "process"):
+                    valid = True
+                    logger.info("MatAnyone model validated")
+                    
+            return valid
+            
         except Exception as e:
             logger.error(f"Model validation failed: {e}")
             return False
 
     def get_model_info(self) -> Dict[str, Any]:
+        """Get detailed information about loaded models"""
         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()
         
-        # Add wrapper status
-        info["matanyone_wrapper_active"] = self._matanyone_wrapper is not None
+        # Add SAM2 info
+        info["sam2"] = self.sam2_loader.get_info() if self.sam2_loader else {}
+        
+        # Add MatAnyone info
+        info["matanyone"] = self.matanyone_loader.get_info() if self.matanyone_loader else {}
         
         return info
 
     def get_load_summary(self) -> str:
+        """Get human-readable loading summary"""
         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"
+            return "No models loaded"
+            
+        lines = []
+        lines.append(f"Models loaded in {self.loading_stats['total_load_time']:.1f}s")
+        
         if self.sam2_predictor:
-            summary += f"✓ SAM2: {sam2_time:.1f}s (ID: {self.sam2_predictor.model_id})\n"
+            lines.append(f"✓ SAM2: {self.loading_stats['sam2_load_time']:.1f}s")
+            lines.append(f"  Model: {self.sam2_predictor.model_id}")
         else:
-            summary += "✗ SAM2: Failed (using fallback)\n"
+            lines.append("✗ SAM2: Failed to load")
+            
         if self.matanyone_model:
-            summary += f"✓ MatAnyOne: {matanyone_time:.1f}s (ID: {self.matanyone_model.model_id})\n"
-            if self._matanyone_wrapper:
-                summary += "  └─ Enhanced wrapper active\n"
+            lines.append(f"✓ MatAnyone: {self.loading_stats['matanyone_load_time']:.1f}s")
+            lines.append(f"  Model: {self.matanyone_model.model_id}")
         else:
-            summary += "✗ MatAnyOne: Failed (using simple refinement)\n"
-        summary += f"Device: {self.device}"
-        return summary
+            lines.append("✗ MatAnyone: Failed to load")
+            
+        lines.append(f"Device: {self.device}")
+        
+        return "\n".join(lines)
 
     def cleanup(self):
+        """Clean up all resources"""
         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.
-        Also: install a shape guard so every call is safe (no 5D tensors).
-        """
-        from matanyone import InferenceCore
-
-        t0 = time.time()
-        processor = InferenceCore("PeiqingYang/MatAnyone")
-
-        # ------------------- BEGIN: GLOBAL SHAPE GUARD PATCH -------------------
-        try:
-            # Lazy import coercers; provide minimal fallbacks if missing.
-            try:
-                from utils.interop import (
-                    ensure_image_nchw,
-                    ensure_mask_for_matanyone,
-                    log_shape,
-                )
-            except Exception as imp_err:
-                logger.warning(f"utils.interop not available ({imp_err}); using minimal inline coercers")
-
-                def log_shape(tag: str, t: torch.Tensor) -> None:
-                    try:
-                        mn = float(t.min()) if t.numel() else float("nan")
-                        mx = float(t.max()) if t.numel() else float("nan")
-                        print(f"[MatAny.guard] {tag}: shape={tuple(t.shape)} dtype={t.dtype} device={t.device} "
-                              f"range=[{mn:.4f},{mx:.4f}]")
-                    except Exception:
-                        pass
-
-                def _to_float01(x: torch.Tensor) -> torch.Tensor:
-                    x = x.to(torch.float32)
-                    if x.max() > 1.0:
-                        x = x / 255.0
-                    return x.clamp_(0.0, 1.0)
-
-                def _squeeze_bt(x: torch.Tensor) -> torch.Tensor:
-                    if x.ndim == 5:
-                        # (B,T,C,H,W) → drop T if 1
-                        if x.shape[1] == 1:
-                            x = x.squeeze(1)
-                        if x.ndim == 5 and x.shape[0] == 1:
-                            x = x.squeeze(0)
-                    if x.ndim == 4 and x.shape[0] == 1 and x.shape[1] == 1 and x.shape[-3] == 3:
-                        x = x.squeeze(1)
-                    return x
-
-                def ensure_image_nchw(img: torch.Tensor, device=self.device, want_batched: bool = True) -> torch.Tensor:
-                    img = img.to(device)
-                    img = _squeeze_bt(img)
-                    if img.ndim == 3:
-                        # CHW or HWC
-                        if img.shape[0] in (1, 3):
-                            chw = img
-                        else:
-                            chw = img.permute(2, 0, 1)
-                        chw = _to_float01(chw.contiguous())
-                        return chw.unsqueeze(0) if want_batched else chw
-                    if img.ndim == 4:
-                        N, A, B, C = img.shape
-                        if A == 3:
-                            nchw = img
-                        elif C == 3:
-                            nchw = img.permute(0, 3, 1, 2)
-                        else:
-                            raise AssertionError(f"Cannot infer channels in image: {tuple(img.shape)}")
-                        nchw = _to_float01(nchw.contiguous())
-                        return nchw if want_batched else nchw[0]
-                    raise AssertionError(f"Bad image dims: {tuple(img.shape)}")
-
-                def ensure_mask_for_matanyone(mask: torch.Tensor, *, idx_mask: bool = False,
-                                              threshold: float = 0.5, keep_soft: bool = False,
-                                              device=self.device) -> torch.Tensor:
-                    mask = mask.to(device)
-                    mask = _squeeze_bt(mask)
-                    if idx_mask:
-                        if mask.ndim == 3:
-                            if mask.shape[0] == 1:
-                                idx = (mask[0] >= threshold).to(torch.long)
-                            else:
-                                idx = torch.argmax(mask, dim=0).to(torch.long)
-                                idx = (idx > 0).to(torch.long)
-                        elif mask.ndim == 2:
-                            idx = (mask >= threshold).to(torch.long)
-                        else:
-                            raise AssertionError(f"idx mask must be 2D or 3D; got {tuple(mask.shape)}")
-                        return idx
-                    # channel mask
-                    if mask.ndim == 2:
-                        out = mask.unsqueeze(0)
-                    elif mask.ndim == 3:
-                        if mask.shape[0] == 1:
-                            out = mask
-                        else:
-                            areas = mask.sum(dim=(-2, -1))
-                            out = mask[areas.argmax():areas.argmax()+1]
-                    else:
-                        raise AssertionError(f"mask must be 2D/3D; got {tuple(mask.shape)}")
-                    out = out.to(torch.float32)
-                    if not keep_soft:
-                        out = (out >= threshold).to(torch.float32)
-                    return out.clamp_(0.0, 1.0).contiguous()
-
-            def _guarded_factory(core_obj, method_name: str):
-                core_step = getattr(core_obj, method_name)
-
-                def wrapped_step(*args, **kwargs):
-                    # Extract image/mask/idx_mask whether passed positionally or by name
-                    image = kwargs.get("image", None)
-                    mask = kwargs.get("mask", None)
-                    idx_mask = kwargs.get("idx_mask", kwargs.get("index_mask", False))
-
-                    # Positional fallback guess: (image, mask, ...)
-                    if image is None and len(args) >= 1:
-                        image = args[0]
-                    if mask is None and len(args) >= 2:
-                        mask = args[1]
-
-                    # Coerce shapes
-                    img_nchw = ensure_image_nchw(image, device=self.device, want_batched=True)
-                    log_shape("image_nchw", img_nchw)
-
-                    if idx_mask:
-                        m_fixed = ensure_mask_for_matanyone(mask, idx_mask=True, device=img_nchw.device)
-                        log_shape("idx_hw", m_fixed)
-                    else:
-                        m_fixed = ensure_mask_for_matanyone(mask, idx_mask=False, threshold=0.5, keep_soft=False, device=img_nchw.device)
-                        log_shape("mask_c_hw", m_fixed)
-
-                    # Rebuild kwargs without duplicates
-                    new_kwargs = dict(kwargs)
-                    new_kwargs["idx_mask"] = bool(idx_mask)
-                    new_kwargs["image"] = img_nchw[0]  # common: CHW image
-
-                    if idx_mask:
-                        new_kwargs["mask"] = m_fixed  # (H,W) long
-                    else:
-                        new_kwargs["mask"] = m_fixed  # (1,H,W) float
-
-                    # Try unbatched first, then batched fallback if needed
-                    try:
-                        return core_step(**new_kwargs)
-                    except Exception as e1:
-                        logger.debug(f"MatAnyOne step (CHW) failed, retrying batched NCHW: {e1}")
-                        new_kwargs["image"] = img_nchw  # (1,3,H,W)
-                        try:
-                            return core_step(**new_kwargs)
-                        except Exception as e2:
-                            logger.error(f"MatAnyOne guarded call failed (both modes). Last error: {e2}")
-                            raise
-
-                return wrapped_step
-
-            if hasattr(processor, "step"):
-                processor.step = _guarded_factory(processor, "step")
-                logger.info("Patched MatAnyOne InferenceCore.step with shape guard")
-            if hasattr(processor, "process"):
-                processor.process = _guarded_factory(processor, "process")
-                logger.info("Patched MatAnyOne InferenceCore.process with shape guard")
-        except Exception as guard_err:
-            logger.warning(f"Could not install MatAnyOne guard: {guard_err}")
-        # -------------------- END: GLOBAL SHAPE GUARD PATCH --------------------
-
-        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, **kwargs):
-                # 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:
+        """Internal cleanup of loaded models"""
+        # Clean up SAM2
+        if self.sam2_loader:
+            self.sam2_loader.cleanup()
+        if self.sam2_predictor:
             del self.sam2_predictor
             self.sam2_predictor = None
-        if self.matanyone_model is not None:
+            
+        # Clean up MatAnyone
+        if self.matanyone_loader:
+            self.matanyone_loader.cleanup()
+        if self.matanyone_model:
             del self.matanyone_model
             self.matanyone_model = None
-        if self._matanyone_wrapper is not None:
-            del self._matanyone_wrapper
-            self._matanyone_wrapper = None
+            
+        # Clear CUDA cache
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
+            
+        # Garbage collection
         gc.collect()
+        
         logger.debug("Model cleanup completed")