model_loader.py

"""
Model Loading Module
Handles loading and validation of SAM2 and MatAnyone AI models
"""

# ============================================================================ #
# IMPORTS AND DEPENDENCIES
# ============================================================================ #

import os
import gc
import sys
import time
import shutil
import logging
import tempfile
import traceback
from typing import Optional, Dict, Any, Tuple, Union
from pathlib import Path

import torch
import gradio as gr
from omegaconf import DictConfig, OmegaConf

# Import modular components
import exceptions
import device_manager
import memory_manager

logger = logging.getLogger(__name__)

# ============================================================================ #
# HARD CACHE CLEANER
# ============================================================================ #

class HardCacheCleaner:
    """
    Comprehensive cache cleaning system to resolve SAM2 loading issues
    Clears Python module cache, HuggingFace cache, and temp files
    """
    
    @staticmethod
    def clean_all_caches(verbose: bool = True):
        """Clean all caches that might interfere with SAM2 loading"""
        
        if verbose:
            logger.info("Starting comprehensive cache cleanup...")
        
        # 1. Clean Python module cache
        HardCacheCleaner._clean_python_cache(verbose)
        
        # 2. Clean HuggingFace cache
        HardCacheCleaner._clean_huggingface_cache(verbose)
        
        # 3. Clean PyTorch cache
        HardCacheCleaner._clean_pytorch_cache(verbose)
        
        # 4. Clean temp directories
        HardCacheCleaner._clean_temp_directories(verbose)
        
        # 5. Clear import cache
        HardCacheCleaner._clear_import_cache(verbose)
        
        # 6. Force garbage collection
        HardCacheCleaner._force_gc_cleanup(verbose)
        
        if verbose:
            logger.info("Cache cleanup completed")
    
    @staticmethod
    def _clean_python_cache(verbose: bool = True):
        """Clean Python bytecode cache"""
        try:
            # Clear sys.modules cache for SAM2 related modules
            sam2_modules = [key for key in sys.modules.keys() if 'sam2' in key.lower()]
            for module in sam2_modules:
                if verbose:
                    logger.info(f"Removing cached module: {module}")
                del sys.modules[module]
            
            # Clear __pycache__ directories
            for root, dirs, files in os.walk("."):
                for dir_name in dirs[:]:  # Use slice to modify list during iteration
                    if dir_name == "__pycache__":
                        cache_path = os.path.join(root, dir_name)
                        if verbose:
                            logger.info(f"Removing __pycache__: {cache_path}")
                        shutil.rmtree(cache_path, ignore_errors=True)
                        dirs.remove(dir_name)
            
        except Exception as e:
            logger.warning(f"Python cache cleanup failed: {e}")
    
    @staticmethod
    def _clean_huggingface_cache(verbose: bool = True):
        """Clean HuggingFace model cache"""
        try:
            cache_paths = [
                os.path.expanduser("~/.cache/huggingface/"),
                os.path.expanduser("~/.cache/torch/"),
                "./checkpoints/",
                "./.cache/",
            ]
            
            for cache_path in cache_paths:
                if os.path.exists(cache_path):
                    if verbose:
                        logger.info(f"Cleaning cache directory: {cache_path}")
                    
                    # Remove SAM2 specific files
                    for root, dirs, files in os.walk(cache_path):
                        for file in files:
                            if any(pattern in file.lower() for pattern in ['sam2', 'segment-anything-2']):
                                file_path = os.path.join(root, file)
                                try:
                                    os.remove(file_path)
                                    if verbose:
                                        logger.info(f"Removed cached file: {file_path}")
                                except:
                                    pass
                        
                        for dir_name in dirs[:]:
                            if any(pattern in dir_name.lower() for pattern in ['sam2', 'segment-anything-2']):
                                dir_path = os.path.join(root, dir_name)
                                try:
                                    shutil.rmtree(dir_path, ignore_errors=True)
                                    if verbose:
                                        logger.info(f"Removed cached directory: {dir_path}")
                                    dirs.remove(dir_name)
                                except:
                                    pass
                                    
        except Exception as e:
            logger.warning(f"HuggingFace cache cleanup failed: {e}")
    
    @staticmethod
    def _clean_pytorch_cache(verbose: bool = True):
        """Clean PyTorch cache"""
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                if verbose:
                    logger.info("Cleared PyTorch CUDA cache")
        except Exception as e:
            logger.warning(f"PyTorch cache cleanup failed: {e}")
    
    @staticmethod
    def _clean_temp_directories(verbose: bool = True):
        """Clean temporary directories"""
        try:
            temp_dirs = [tempfile.gettempdir(), "/tmp", "./tmp", "./temp"]
            
            for temp_dir in temp_dirs:
                if os.path.exists(temp_dir):
                    for item in os.listdir(temp_dir):
                        if 'sam2' in item.lower() or 'segment' in item.lower():
                            item_path = os.path.join(temp_dir, item)
                            try:
                                if os.path.isfile(item_path):
                                    os.remove(item_path)
                                elif os.path.isdir(item_path):
                                    shutil.rmtree(item_path, ignore_errors=True)
                                if verbose:
                                    logger.info(f"Removed temp item: {item_path}")
                            except:
                                pass
                                
        except Exception as e:
            logger.warning(f"Temp directory cleanup failed: {e}")
    
    @staticmethod
    def _clear_import_cache(verbose: bool = True):
        """Clear Python import cache"""
        try:
            import importlib
            
            # Invalidate import caches
            importlib.invalidate_caches()
            
            if verbose:
                logger.info("Cleared Python import cache")
                
        except Exception as e:
            logger.warning(f"Import cache cleanup failed: {e}")
    
    @staticmethod
    def _force_gc_cleanup(verbose: bool = True):
        """Force garbage collection"""
        try:
            collected = gc.collect()
            if verbose:
                logger.info(f"Garbage collection freed {collected} objects")
        except Exception as e:
            logger.warning(f"Garbage collection failed: {e}")

# ============================================================================ #
# MODEL LOADER CLASS - MAIN INTERFACE
# ============================================================================ #

class ModelLoader:
    """
    Comprehensive model loading and management for SAM2 and MatAnyone
    Handles automatic config detection, multiple fallback strategies, and memory management
    """
    
    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
        self.matanyone_model = None
        self.matanyone_core = None
        
        # Configuration paths
        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()

# ============================================================================ #
# INITIALIZATION AND SETUP
# ============================================================================ #
    
    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}")

# ============================================================================ #
# MAIN MODEL LOADING ORCHESTRATION
# ============================================================================ #

    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_callback: Progress update callback
            cancel_event: Event to check for cancellation
            
        Returns:
            Tuple of (sam2_predictor, 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...")
            
            # Clear any existing models
            self._cleanup_models()
            
            # Load SAM2 first (typically faster)
            logger.info("Loading SAM2 predictor...")
            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 exceptions.ModelLoadingError("SAM2", "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...")
            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_callback)
            
            if self.matanyone_model is None:
                raise exceptions.ModelLoadingError("MatAnyone", "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:.1f}s")
            
            # Final setup
            total_time = time.time() - start_time
            self.loading_stats['total_load_time'] = total_time
            self.loading_stats['models_loaded'] = True
            
            if progress_callback:
                progress_callback(1.0, "Models loaded successfully!")
                
            logger.info(f"All models loaded successfully 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()}")
            
            # Cleanup on failure
            self._cleanup_models()
            self.loading_stats['models_loaded'] = False
            
            if progress_callback:
                progress_callback(1.0, f"Error: {error_msg}")
            
            return None, None

# ============================================================================ #
# SAM2 MODEL LOADING - HUGGINGFACE TRANSFORMERS APPROACH
# ============================================================================ #

    def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
        """
        Load SAM2 using HuggingFace Transformers integration with cache cleanup
        This method works reliably on HuggingFace Spaces without config file issues
        
        Args:
            progress_callback: Progress update callback
            
        Returns:
            SAM2 model or None
        """
        logger.info("=== USING NEW HF TRANSFORMERS SAM2 LOADER ===")
        
        # Step 1: Clean caches before loading
        if progress_callback:
            progress_callback(0.15, "Cleaning caches...")
        
        HardCacheCleaner.clean_all_caches(verbose=True)
        
        # Step 2: Determine model size based on device memory
        model_size = "large"  # default
        if hasattr(self.device_manager, 'get_device_memory_gb'):
            try:
                memory_gb = self.device_manager.get_device_memory_gb()
                if memory_gb < 4:
                    model_size = "tiny"
                elif memory_gb < 8:
                    model_size = "base"
                logger.info(f"Selected SAM2 {model_size} based on {memory_gb}GB memory")
            except Exception as e:
                logger.warning(f"Could not determine device memory: {e}")
        
        # Step 3: Try multiple HuggingFace approaches
        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["large"])
        
        if progress_callback:
            progress_callback(0.3, f"Loading SAM2 {model_size}...")
        
        # Method 1: HuggingFace Transformers Pipeline (most reliable)
        try:
            logger.info("Trying Transformers pipeline approach...")
            from transformers import pipeline
            
            sam2_pipeline = pipeline(
                "mask-generation",
                model=model_id,
                device=0 if str(self.device) == "cuda" else -1
            )
            
            logger.info("SAM2 loaded successfully via Transformers pipeline")
            return sam2_pipeline
            
        except Exception as e:
            logger.warning(f"Pipeline approach failed: {e}")
        
        # Method 2: Direct Transformers classes
        try:
            logger.info("Trying direct Transformers classes...")
            from transformers import Sam2Processor, Sam2Model
            
            processor = Sam2Processor.from_pretrained(model_id)
            model = Sam2Model.from_pretrained(model_id).to(self.device)
            
            logger.info("SAM2 loaded successfully via Transformers classes")
            return {"model": model, "processor": processor}
            
        except Exception as e:
            logger.warning(f"Direct class approach failed: {e}")
        
        # Method 3: Official SAM2 with from_pretrained
        try:
            logger.info("Trying official SAM2 from_pretrained...")
            from sam2.sam2_image_predictor import SAM2ImagePredictor
            
            predictor = SAM2ImagePredictor.from_pretrained(model_id)
            
            logger.info("SAM2 loaded successfully via official from_pretrained")
            return predictor
            
        except Exception as e:
            logger.warning(f"Official from_pretrained approach failed: {e}")
        
        # Method 4: Fallback to direct checkpoint download
        try:
            logger.info("Trying fallback checkpoint approach...")
            from huggingface_hub import hf_hub_download
            from transformers import Sam2Model
            
            # Download checkpoint directly
            checkpoint_path = hf_hub_download(
                repo_id=model_id,
                filename="model.safetensors" if "sam2.1" in model_id else "pytorch_model.bin"
            )
            
            logger.info(f"Downloaded checkpoint to: {checkpoint_path}")
            
            # Load with minimal approach
            model = Sam2Model.from_pretrained(model_id)
            model = model.to(self.device)
            
            logger.info("SAM2 loaded successfully via fallback approach")
            return model
            
        except Exception as e:
            logger.warning(f"Fallback approach failed: {e}")
        
        logger.error("All SAM2 loading methods failed")
        return None

# ============================================================================ #
# MATANYONE MODEL LOADING - MULTIPLE STRATEGIES
# ============================================================================ #

    def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
        """
        Load MatAnyone model with multiple import strategies
        
        Args:
            progress_callback: Progress update 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}...")
                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:
                    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
        
        logger.error("All MatAnyone loading strategies failed")
        return None, None

# ============================================================================ #
# MATANYONE LOADING STRATEGIES
# ============================================================================ #

    def _load_matanyone_strategy_1(self):
        """MatAnyone loading strategy 1: Official HuggingFace InferenceCore"""
        from matanyone import InferenceCore
        
        # Initialize with the official model repo
        processor = InferenceCore("PeiqingYang/MatAnyone")
        return processor, processor
    
    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: Direct model class"""
        from matanyone.model.matanyone import MatAnyone
        
        model = MatAnyone.from_pretrained("not-lain/matanyone")
        return model, model

# ============================================================================ #
# MODEL MANAGEMENT AND CLEANUP
# ============================================================================ #

    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_aggressive()
        gc.collect()
        
        logger.debug("Model cleanup completed")
    
    def cleanup(self):
        """Clean up all resources"""
        self._cleanup_models()
        logger.info("ModelLoader cleanup completed")

# ============================================================================ #
# MODEL INFORMATION AND STATUS
# ============================================================================ #

    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).__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 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

# ============================================================================ #
# MODEL VALIDATION AND TESTING
# ============================================================================ #

    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

# ============================================================================ #
# UTILITY METHODS
# ============================================================================ #

    def reload_models(self, progress_callback: Optional[callable] = None) -> Tuple[Any, Any]:
        """
        Reload all models (useful for error recovery)
        
        Args:
            progress_callback: Progress update callback
            
        Returns:
            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_callback)
    
    @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
        )