Update model_loader.py

model_loader.py

@@ -17,9 +17,10 @@
 import gradio as gr
 from omegaconf import DictConfig, OmegaConf
 
-from exceptions import ModelLoadingError, ConfigurationError
-from device_manager import DeviceManager
-from memory_manager import MemoryManager
+# Import modular components
+import exceptions
+import device_manager
+import memory_manager
 
 logger = logging.getLogger(__name__)
 
@@ -28,10 +29,10 @@ class ModelLoader:
     Comprehensive model loading and management for SAM2 and MatAnyone
     """
     
-    def __init__(self, device_manager: DeviceManager, memory_manager: MemoryManager):
-        self.device_manager = device_manager
-        self.memory_manager = memory_manager
-        self.device = device
+    def __init__(self, device_mgr: device_manager.DeviceManager, memory_mgr: memory_manager.MemoryManager):
+        self.device_manager = device_mgr
+        self.memory_manager = memory_mgr
+        self.device = self.device_manager.get_optimal_device()
         
         # Model storage
         self.sam2_predictor = None
@@ -75,33 +76,37 @@ def patched_get_config(self):
         except (ImportError, AttributeError) as e:
             logger.warning(f"Could not apply Gradio monkey patch: {e}")
     
-    def load_all_models(self, progress: Optional[gr.Progress] = None) -> bool:
+    def load_all_models(self, progress_callback: Optional[callable] = None, cancel_event=None) -> Tuple[Any, Any]:
         """
         Load both SAM2 and MatAnyone models with comprehensive error handling
         
         Args:
-            progress: Gradio progress callback
+            progress_callback: Progress update callback
+            cancel_event: Event to check for cancellation
             
         Returns:
-            bool: True if all models loaded successfully
+            Tuple of (sam2_predictor, matanyone_model)
         """
         start_time = time.time()
         self.loading_stats['loading_attempts'] += 1
         
         try:
             logger.info("Starting model loading process...")
-            self._maybe_progress(progress, 0.0, "Initializing model loading...")
+            if progress_callback:
+                progress_callback(0.0, "Initializing model loading...")
             
             # Clear any existing models
             self._cleanup_models()
             
             # Load SAM2 first (typically faster)
             logger.info("Loading SAM2 predictor...")
-            self._maybe_progress(progress, 0.1, "Loading SAM2 predictor...")
-            self.sam2_predictor = self._load_sam2_predictor(progress)
+            if progress_callback:
+                progress_callback(0.1, "Loading SAM2 predictor...")
+                
+            self.sam2_predictor = self._load_sam2_predictor(progress_callback)
             
             if self.sam2_predictor is None:
-                raise ModelLoadingError("Failed to load SAM2 predictor")
+                raise exceptions.ModelLoadingError("Failed to load SAM2 predictor")
             
             sam2_time = time.time() - start_time
             self.loading_stats['sam2_load_time'] = sam2_time
@@ -109,13 +114,15 @@ def load_all_models(self, progress: Optional[gr.Progress] = None) -> bool:
             
             # Load MatAnyone
             logger.info("Loading MatAnyone model...")
-            self._maybe_progress(progress, 0.6, "Loading MatAnyone model...")
+            if progress_callback:
+                progress_callback(0.6, "Loading MatAnyone model...")
+                
             matanyone_start = time.time()
             
-            self.matanyone_model, self.matanyone_core = self._load_matanyone_model(progress)
+            self.matanyone_model, self.matanyone_core = self._load_matanyone_model(progress_callback)
             
             if self.matanyone_model is None:
-                raise ModelLoadingError("Failed to load MatAnyone model")
+                raise exceptions.ModelLoadingError("Failed to load MatAnyone model")
             
             matanyone_time = time.time() - matanyone_start
             self.loading_stats['matanyone_load_time'] = matanyone_time
@@ -126,10 +133,12 @@ def load_all_models(self, progress: Optional[gr.Progress] = None) -> bool:
             self.loading_stats['total_load_time'] = total_time
             self.loading_stats['models_loaded'] = True
             
-            self._maybe_progress(progress, 1.0, "Models loaded successfully!")
+            if progress_callback:
+                progress_callback(1.0, "Models loaded successfully!")
+                
             logger.info(f"All models loaded successfully in {total_time:.2f}s")
             
-            return True
+            return self.sam2_predictor, self.matanyone_model
             
         except Exception as e:
             error_msg = f"Model loading failed: {str(e)}"
@@ -139,23 +148,23 @@ def load_all_models(self, progress: Optional[gr.Progress] = None) -> bool:
             self._cleanup_models()
             self.loading_stats['models_loaded'] = False
             
-            if progress:
-                progress(1.0, desc=f"Error: {error_msg}")
+            if progress_callback:
+                progress_callback(1.0, f"Error: {error_msg}")
             
-            raise ModelLoadingError(error_msg) from e
+            return None, None
     
-    def _load_sam2_predictor(self, progress: Optional[gr.Progress] = None):
+    def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
         """
         Load SAM2 predictor with multiple fallback strategies
         
         Args:
-            progress: Gradio progress callback
+            progress_callback: Progress update callback
             
         Returns:
             SAM2ImagePredictor or None
         """
         if not os.path.isdir(self.configs_dir):
-            raise ModelLoadingError(f"SAM2 Configs directory not found at '{self.configs_dir}'")
+            logger.warning(f"SAM2 Configs directory not found at '{self.configs_dir}', trying fallback loading")
         
         def try_load_sam2(config_name_with_yaml: str, checkpoint_name: str):
             """Attempt to load SAM2 with given config and checkpoint"""
@@ -166,31 +175,38 @@ def try_load_sam2(config_name_with_yaml: str, checkpoint_name: str):
                 # Download checkpoint if needed
                 if not os.path.exists(checkpoint_path):
                     logger.info(f"Downloading {checkpoint_name} from Hugging Face Hub...")
-                    self._maybe_progress(progress, 0.2, f"Downloading {checkpoint_name}...")
+                    if progress_callback:
+                        progress_callback(0.2, f"Downloading {checkpoint_name}...")
                     
-                    from huggingface_hub import hf_hub_download
-                    repo = f"facebook/{config_name_with_yaml.replace('.yaml','')}"
-                    checkpoint_path = hf_hub_download(
-                        repo_id=repo,
-                        filename=checkpoint_name,
-                        cache_dir=self.checkpoints_dir,
-                        local_dir_use_symlinks=False
-                    )
-                    logger.info(f"Download complete: {checkpoint_path}")
+                    try:
+                        from huggingface_hub import hf_hub_download
+                        repo = f"facebook/{config_name_with_yaml.replace('.yaml','')}"
+                        checkpoint_path = hf_hub_download(
+                            repo_id=repo,
+                            filename=checkpoint_name,
+                            cache_dir=self.checkpoints_dir,
+                            local_dir_use_symlinks=False
+                        )
+                        logger.info(f"Download complete: {checkpoint_path}")
+                    except Exception as download_error:
+                        logger.warning(f"Failed to download {checkpoint_name}: {download_error}")
+                        return None
                 
-                # Reset and initialize Hydra
-                if hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
-                    hydra.core.global_hydra.GlobalHydra.instance().clear()
-                
-                hydra.initialize(
-                    version_base=None,
-                    config_path=os.path.relpath(self.configs_dir),
-                    job_name=f"sam2_load_{int(time.time())}"
-                )
+                # Reset and initialize Hydra if configs directory exists
+                if os.path.isdir(self.configs_dir):
+                    if hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
+                        hydra.core.global_hydra.GlobalHydra.instance().clear()
+                    
+                    hydra.initialize(
+                        version_base=None,
+                        config_path=os.path.relpath(self.configs_dir),
+                        job_name=f"sam2_load_{int(time.time())}"
+                    )
                 
                 # Build SAM2 model
                 config_name = config_name_with_yaml.replace(".yaml", "")
-                self._maybe_progress(progress, 0.4, f"Building {config_name}...")
+                if progress_callback:
+                    progress_callback(0.4, f"Building {config_name}...")
                 
                 from sam2.build_sam import build_sam2
                 from sam2.sam2_image_predictor import SAM2ImagePredictor
@@ -217,25 +233,29 @@ def try_load_sam2(config_name_with_yaml: str, checkpoint_name: str):
         
         # Prioritize model size based on device memory
         if hasattr(self.device_manager, 'get_device_memory_gb'):
-            memory_gb = self.device_manager.get_device_memory_gb()
-            if memory_gb < 4:
-                model_attempts = model_attempts[2:]  # Only tiny and small
-            elif memory_gb < 8:
-                model_attempts = model_attempts[1:]  # Skip large
+            try:
+                memory_gb = self.device_manager.get_device_memory_gb()
+                if memory_gb < 4:
+                    model_attempts = model_attempts[2:]  # Only tiny and small
+                elif memory_gb < 8:
+                    model_attempts = model_attempts[1:]  # Skip large
+            except Exception as e:
+                logger.warning(f"Could not determine device memory: {e}")
         
         for config_yaml, checkpoint_pt in model_attempts:
             predictor = try_load_sam2(config_yaml, checkpoint_pt)
             if predictor is not None:
                 return predictor
         
-        raise ModelLoadingError("All SAM2 model loading attempts failed")
+        logger.error("All SAM2 model loading attempts failed")
+        return None
     
-    def _load_matanyone_model(self, progress: Optional[gr.Progress] = None):
+    def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
         """
         Load MatAnyone model with multiple import strategies
         
         Args:
-            progress: Gradio progress callback
+            progress_callback: Progress update callback
             
         Returns:
             Tuple[model, core] or (None, None)
@@ -250,7 +270,8 @@ def _load_matanyone_model(self, progress: Optional[gr.Progress] = None):
         for i, strategy in enumerate(import_strategies, 1):
             try:
                 logger.info(f"Trying MatAnyone loading strategy {i}...")
-                self._maybe_progress(progress, 0.7 + (i * 0.05), f"MatAnyone strategy {i}...")
+                if progress_callback:
+                    progress_callback(0.7 + (i * 0.05), f"MatAnyone strategy {i}...")
                 
                 model, core = strategy()
                 if model is not None and core is not None:
@@ -261,7 +282,8 @@ def _load_matanyone_model(self, progress: Optional[gr.Progress] = None):
                 logger.warning(f"MatAnyone strategy {i} failed: {e}")
                 continue
         
-        raise ModelLoadingError("All MatAnyone loading strategies failed")
+        logger.error("All MatAnyone loading strategies failed")
+        return None, None
     
     def _load_matanyone_strategy_1(self):
         """MatAnyone loading strategy 1: Direct model import"""
@@ -350,14 +372,6 @@ def _cleanup_models(self):
         
         logger.debug("Model cleanup completed")
     
-    def _maybe_progress(self, progress: Optional[gr.Progress], value: float, desc: str):
-        """Update progress if callback is available"""
-        if progress is not None:
-            try:
-                progress(value, desc=desc)
-            except Exception as e:
-                logger.debug(f"Progress update failed: {e}")
-    
     def get_model_info(self) -> Dict[str, Any]:
         """
         Get information about loaded models
@@ -387,6 +401,26 @@ def get_model_info(self) -> Dict[str, Any]:
         
         return info
     
+    def get_status(self) -> Dict[str, Any]:
+        """Get model loader status for backward compatibility"""
+        return self.get_model_info()
+    
+    def get_load_summary(self) -> str:
+        """Get a human-readable summary of model loading"""
+        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 successfully in {total_time:.1f}s\n"
+        summary += f"SAM2: {sam2_time:.1f}s\n"
+        summary += f"MatAnyone: {matanyone_time:.1f}s\n"
+        summary += f"Device: {self.device}"
+        
+        return summary
+    
     def validate_models(self) -> bool:
         """
         Validate that models are properly loaded and functional
@@ -411,21 +445,26 @@ def validate_models(self) -> bool:
             logger.error(f"Model validation failed: {e}")
             return False
     
-    def reload_models(self, progress: Optional[gr.Progress] = None) -> bool:
+    def reload_models(self, progress_callback: Optional[callable] = None) -> Tuple[Any, Any]:
         """
         Reload all models (useful for error recovery)
         
         Args:
-            progress: Gradio progress callback
+            progress_callback: Progress update callback
             
         Returns:
-            bool: True if reload successful
+            Tuple of (sam2_predictor, matanyone_model)
         """
         logger.info("Reloading models...")
         self._cleanup_models()
         self.loading_stats['models_loaded'] = False
         
-        return self.load_all_models(progress)
+        return self.load_all_models(progress_callback)
+    
+    def cleanup(self):
+        """Clean up all resources"""
+        self._cleanup_models()
+        logger.info("ModelLoader cleanup completed")
     
     @property
     def models_ready(self) -> bool:
@@ -434,4 +473,4 @@ def models_ready(self) -> bool:
             self.loading_stats['models_loaded'] and
             self.sam2_predictor is not None and
             self.matanyone_model is not None
-        ) 
+        )