Create model_loader.py

model_loader.py

@@ -0,0 +1,437 @@
+"""
+Model Loading Module
+Handles loading and validation of SAM2 and MatAnyone AI models
+"""
+
+import os
+import gc
+import time
+import logging
+import tempfile
+import traceback
+from typing import Optional, Dict, Any, Tuple, Union
+from pathlib import Path
+
+import torch
+import hydra
+import gradio as gr
+from omegaconf import DictConfig, OmegaConf
+
+from exceptions import ModelLoadingError, ConfigurationError
+from device_manager import DeviceManager
+from memory_manager import MemoryManager
+
+logger = logging.getLogger(__name__)
+
+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_manager.get_optimal_device()
+        
+        # Model storage
+        self.sam2_predictor = None
+        self.matanyone_model = None
+        self.matanyone_core = None
+        
+        # Configuration paths
+        self.configs_dir = os.path.abspath("Configs")
+        self.checkpoints_dir = "./checkpoints"
+        os.makedirs(self.checkpoints_dir, exist_ok=True)
+        
+        # Model loading statistics
+        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}")
+        self._apply_gradio_patch()
+    
+    def _apply_gradio_patch(self):
+        """Apply Gradio schema monkey patch to prevent validation errors"""
+        try:
+            import gradio.components.base
+            original_get_config = gradio.components.base.Component.get_config
+            
+            def patched_get_config(self):
+                config = original_get_config(self)
+                # Remove problematic keys that cause validation errors
+                config.pop("show_progress_bar", None)
+                config.pop("min_width", None)
+                config.pop("scale", None)
+                return config
+            
+            gradio.components.base.Component.get_config = patched_get_config
+            logger.debug("Applied Gradio schema monkey patch")
+            
+        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:
+        """
+        Load both SAM2 and MatAnyone models with comprehensive error handling
+        
+        Args:
+            progress: Gradio progress callback
+            
+        Returns:
+            bool: True if all models loaded successfully
+        """
+        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...")
+            
+            # 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 self.sam2_predictor is None:
+                raise ModelLoadingError("Failed to load SAM2 predictor")
+            
+            sam2_time = time.time() - start_time
+            self.loading_stats['sam2_load_time'] = sam2_time
+            logger.info(f"SAM2 loaded in {sam2_time:.2f}s")
+            
+            # Load MatAnyone
+            logger.info("Loading MatAnyone model...")
+            self._maybe_progress(progress, 0.6, "Loading MatAnyone model...")
+            matanyone_start = time.time()
+            
+            self.matanyone_model, self.matanyone_core = self._load_matanyone_model(progress)
+            
+            if self.matanyone_model is None:
+                raise ModelLoadingError("Failed to load MatAnyone model")
+            
+            matanyone_time = time.time() - matanyone_start
+            self.loading_stats['matanyone_load_time'] = matanyone_time
+            logger.info(f"MatAnyone loaded in {matanyone_time:.2f}s")
+            
+            # Final setup
+            total_time = time.time() - start_time
+            self.loading_stats['total_load_time'] = total_time
+            self.loading_stats['models_loaded'] = True
+            
+            self._maybe_progress(progress, 1.0, "Models loaded successfully!")
+            logger.info(f"All models loaded successfully in {total_time:.2f}s")
+            
+            return True
+            
+        except Exception as e:
+            error_msg = f"Model loading failed: {str(e)}"
+            logger.error(f"{error_msg}\n{traceback.format_exc()}")
+            
+            # Cleanup on failure
+            self._cleanup_models()
+            self.loading_stats['models_loaded'] = False
+            
+            if progress:
+                progress(1.0, desc=f"Error: {error_msg}")
+            
+            raise ModelLoadingError(error_msg) from e
+    
+    def _load_sam2_predictor(self, progress: Optional[gr.Progress] = None):
+        """
+        Load SAM2 predictor with multiple fallback strategies
+        
+        Args:
+            progress: Gradio progress 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}'")
+        
+        def try_load_sam2(config_name_with_yaml: str, checkpoint_name: str):
+            """Attempt to load SAM2 with given config and checkpoint"""
+            try:
+                checkpoint_path = os.path.join(self.checkpoints_dir, checkpoint_name)
+                logger.info(f"Attempting SAM2 checkpoint: {checkpoint_path}")
+                
+                # 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}...")
+                    
+                    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}")
+                
+                # 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())}"
+                )
+                
+                # Build SAM2 model
+                config_name = config_name_with_yaml.replace(".yaml", "")
+                self._maybe_progress(progress, 0.4, f"Building {config_name}...")
+                
+                from sam2.build_sam import build_sam2
+                from sam2.sam2_image_predictor import SAM2ImagePredictor
+                
+                sam2_model = build_sam2(config_name, checkpoint_path)
+                sam2_model.to(self.device)
+                predictor = SAM2ImagePredictor(sam2_model)
+                
+                logger.info(f"SAM2 {config_name} loaded successfully on {self.device}")
+                return predictor
+                
+            except Exception as e:
+                error_msg = f"Failed to load SAM2 {config_name_with_yaml}: {e}"
+                logger.warning(error_msg)
+                return None
+        
+        # Try different SAM2 model sizes based on device capabilities
+        model_attempts = [
+            ("sam2_hiera_large.yaml", "sam2_hiera_large.pt"),
+            ("sam2_hiera_base_plus.yaml", "sam2_hiera_base_plus.pt"),
+            ("sam2_hiera_small.yaml", "sam2_hiera_small.pt"),
+            ("sam2_hiera_tiny.yaml", "sam2_hiera_tiny.pt")
+        ]
+        
+        # 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
+        
+        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")
+    
+    def _load_matanyone_model(self, progress: Optional[gr.Progress] = None):
+        """
+        Load MatAnyone model with multiple import strategies
+        
+        Args:
+            progress: Gradio progress callback
+            
+        Returns:
+            Tuple[model, core] or (None, None)
+        """
+        import_strategies = [
+            self._load_matanyone_strategy_1,
+            self._load_matanyone_strategy_2,
+            self._load_matanyone_strategy_3,
+            self._load_matanyone_strategy_4
+        ]
+        
+        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}...")
+                
+                model, core = strategy()
+                if model is not None and core is not None:
+                    logger.info(f"MatAnyone loaded successfully with strategy {i}")
+                    return model, core
+                    
+            except Exception as e:
+                logger.warning(f"MatAnyone strategy {i} failed: {e}")
+                continue
+        
+        raise ModelLoadingError("All MatAnyone loading strategies failed")
+    
+    def _load_matanyone_strategy_1(self):
+        """MatAnyone loading strategy 1: Direct model import"""
+        from matanyone.model.matanyone import MatAnyOne
+        from matanyone.inference.inference_core import InferenceCore
+        
+        cfg = OmegaConf.create({
+            'model': {'name': 'MatAnyOne'},
+            'device': str(self.device),
+            'fp16': True if self.device.type == 'cuda' else False
+        })
+        
+        net = MatAnyOne(cfg)
+        core = InferenceCore(net, cfg)
+        
+        return net, core
+    
+    def _load_matanyone_strategy_2(self):
+        """MatAnyone loading strategy 2: Alternative import paths"""
+        from matanyone import MatAnyOne
+        from matanyone import InferenceCore
+        
+        cfg = OmegaConf.create({
+            'model_name': 'matanyone',
+            'device': str(self.device)
+        })
+        
+        model = MatAnyOne(cfg)
+        core = InferenceCore(model, cfg)
+        
+        return model, core
+    
+    def _load_matanyone_strategy_3(self):
+        """MatAnyone loading strategy 3: Repository-specific imports"""
+        try:
+            from matanyone.models.matanyone import MatAnyOneModel
+            from matanyone.core import InferenceEngine
+        except ImportError:
+            from matanyone.src.models import MatAnyOneModel
+            from matanyone.src.core import InferenceEngine
+        
+        config = {
+            'model_path': None,  # Will use default
+            'device': self.device,
+            'precision': 'fp16' if self.device.type == 'cuda' else 'fp32'
+        }
+        
+        model = MatAnyOneModel.from_pretrained(config)
+        engine = InferenceEngine(model)
+        
+        return model, engine
+    
+    def _load_matanyone_strategy_4(self):
+        """MatAnyone loading strategy 4: Hugging Face Hub approach"""
+        from huggingface_hub import hf_hub_download
+        from matanyone import load_model_from_hub
+        
+        # Try to load from Hugging Face
+        model_path = hf_hub_download(
+            repo_id="PeiqingYang/MatAnyone",
+            filename="pytorch_model.bin",
+            cache_dir=self.checkpoints_dir
+        )
+        
+        model = load_model_from_hub(model_path, device=self.device)
+        
+        return model, model  # Return same object for both
+    
+    def _cleanup_models(self):
+        """Clean up loaded models and free memory"""
+        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 self.matanyone_core is not None:
+            del self.matanyone_core
+            self.matanyone_core = None
+        
+        # Clear GPU cache
+        self.memory_manager.cleanup_gpu_memory()
+        gc.collect()
+        
+        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
+        
+        Returns:
+            Dict with model information and statistics
+        """
+        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:
+            try:
+                info['sam2_model_type'] = type(self.sam2_predictor.model).__name__
+            except:
+                info['sam2_model_type'] = "Unknown"
+        
+        if self.matanyone_model is not None:
+            try:
+                info['matanyone_model_type'] = type(self.matanyone_model).__name__
+            except:
+                info['matanyone_model_type'] = "Unknown"
+        
+        return info
+    
+    def validate_models(self) -> bool:
+        """
+        Validate that models are properly loaded and functional
+        
+        Returns:
+            bool: True if models are valid
+        """
+        try:
+            # Basic validation
+            if not self.loading_stats['models_loaded']:
+                return False
+            
+            if self.sam2_predictor is None or self.matanyone_model is None:
+                return False
+            
+            # Try basic model operations
+            # This could include running a small test inference
+            logger.info("Model validation passed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Model validation failed: {e}")
+            return False
+    
+    def reload_models(self, progress: Optional[gr.Progress] = None) -> bool:
+        """
+        Reload all models (useful for error recovery)
+        
+        Args:
+            progress: Gradio progress callback
+            
+        Returns:
+            bool: True if reload successful
+        """
+        logger.info("Reloading models...")
+        self._cleanup_models()
+        self.loading_stats['models_loaded'] = False
+        
+        return self.load_all_models(progress)
+    
+    @property
+    def models_ready(self) -> bool:
+        """Check if all models are loaded and ready"""
+        return (
+            self.loading_stats['models_loaded'] and
+            self.sam2_predictor is not None and
+            self.matanyone_model is not None
+        )