Delete models/models

models/models/loader.py

@@ -1,491 +0,0 @@
-"""
-Model Loading Module
-Handles loading and validation of SAM2 and MatAnyone AI models
-"""
-
-# ============================================================================
-# IMPORTS AND DEPENDENCIES
-# ============================================================================
-
-import os
-import sys
-import gc
-import traceback
-from pathlib import Path
-from typing import Optional, Dict, Any, Tuple
-import logging
-
-import torch
-import numpy as np
-from huggingface_hub import hf_hub_download
-
-# ============================================================================
-# SYSTEM PATH CONFIGURATION
-# ============================================================================
-
-def setup_paths():
-    """Configure system paths for model imports"""
-    current_dir = Path(__file__).parent.absolute()
-    project_root = current_dir.parent.parent
-    
-    paths_to_add = [
-        str(project_root),
-        str(project_root / "models"),
-        str(project_root / "models" / "sam2"),
-        str(project_root / "models" / "matting"),
-    ]
-    
-    for path in paths_to_add:
-        if path not in sys.path:
-            sys.path.insert(0, path)
-    
-    return project_root
-
-PROJECT_ROOT = setup_paths()
-
-# ============================================================================
-# MODEL IMPORTS (After Path Setup)
-# ============================================================================
-
-try:
-    from sam2.build_sam import build_sam2_video_predictor
-    from models.matting.model_initialization import ModelInitializer
-except ImportError as e:
-    logging.error(f"Failed to import models: {e}")
-    logging.error(f"Current sys.path: {sys.path}")
-    raise
-
-# ============================================================================
-# CONFIGURATION
-# ============================================================================
-
-class ModelConfig:
-    """Model configuration and paths"""
-    
-    # Model identifiers
-    SAM2_MODEL_ID = "facebook/sam2-hiera-large"
-    MATANYONE_REPO = "bytedance/Matting-Anything"
-    
-    # Model filenames
-    SAM2_CHECKPOINT = "sam2_hiera_large.pt"
-    SAM2_CONFIG = "sam2_hiera_l.yaml"
-    MATANYONE_CHECKPOINT = "model_any_mat_vitl.pth"
-    
-    # Default paths
-    DEFAULT_CACHE_DIR = Path.home() / ".cache" / "huggingface" / "hub"
-    
-    # Device configuration
-    CUDA_AVAILABLE = torch.cuda.is_available()
-    DEFAULT_DEVICE = "cuda" if CUDA_AVAILABLE else "cpu"
-    
-    # Memory thresholds (in GB)
-    MIN_GPU_MEMORY = 8.0
-    MIN_RAM = 16.0
-    RECOMMENDED_GPU_MEMORY = 12.0
-    RECOMMENDED_RAM = 32.0
-
-# ============================================================================
-# UTILITY FUNCTIONS
-# ============================================================================
-
-def get_memory_info() -> Dict[str, float]:
-    """Get system memory information"""
-    import psutil
-    
-    memory_info = {
-        'ram_total_gb': psutil.virtual_memory().total / (1024**3),
-        'ram_available_gb': psutil.virtual_memory().available / (1024**3),
-        'ram_used_percent': psutil.virtual_memory().percent
-    }
-    
-    if torch.cuda.is_available():
-        try:
-            gpu_props = torch.cuda.get_device_properties(0)
-            memory_info.update({
-                'gpu_name': gpu_props.name,
-                'gpu_total_gb': gpu_props.total_memory / (1024**3),
-                'gpu_allocated_gb': torch.cuda.memory_allocated(0) / (1024**3),
-                'gpu_reserved_gb': torch.cuda.memory_reserved(0) / (1024**3)
-            })
-        except Exception as e:
-            logging.warning(f"Could not get GPU memory info: {e}")
-    
-    return memory_info
-
-def clean_memory():
-    """Clean up GPU and system memory"""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-    gc.collect()
-
-def format_memory_status(memory_info: Dict[str, float]) -> str:
-    """Format memory information for display"""
-    lines = [
-        "=== System Memory Status ===",
-        f"RAM: {memory_info['ram_available_gb']:.1f}GB / {memory_info['ram_total_gb']:.1f}GB available ({memory_info['ram_used_percent']:.1f}% used)"
-    ]
-    
-    if 'gpu_name' in memory_info:
-        lines.extend([
-            f"GPU: {memory_info['gpu_name']}",
-            f"  Total: {memory_info['gpu_total_gb']:.1f}GB",
-            f"  Allocated: {memory_info['gpu_allocated_gb']:.2f}GB",
-            f"  Reserved: {memory_info['gpu_reserved_gb']:.2f}GB"
-        ])
-    else:
-        lines.append("GPU: Not available")
-    
-    return "\n".join(lines)
-
-# ============================================================================
-# MODEL LOADER CLASS
-# ============================================================================
-
-class ModelLoader:
-    """Manages loading and initialization of AI models"""
-    
-    def __init__(self, cache_dir: Optional[Path] = None, device: Optional[str] = None):
-        """
-        Initialize model loader
-        
-        Args:
-            cache_dir: Directory for caching models
-            device: Device to load models on ('cuda' or 'cpu')
-        """
-        self.cache_dir = Path(cache_dir) if cache_dir else ModelConfig.DEFAULT_CACHE_DIR
-        self.device = device or ModelConfig.DEFAULT_DEVICE
-        
-        # Model instances
-        self.sam2_predictor = None
-        self.matanyone_model = None
-        
-        # Status tracking
-        self.models_loaded = False
-        self.load_errors = []
-        
-        # Setup logging
-        self.logger = logging.getLogger(__name__)
-        
-        # Validate system requirements
-        self._check_system_requirements()
-    
-    def _check_system_requirements(self) -> bool:
-        """Check if system meets minimum requirements"""
-        memory_info = get_memory_info()
-        warnings = []
-        
-        # Check RAM
-        if memory_info['ram_total_gb'] < ModelConfig.MIN_RAM:
-            warnings.append(f"⚠️ Low RAM: {memory_info['ram_total_gb']:.1f}GB (minimum {ModelConfig.MIN_RAM}GB recommended)")
-        
-        # Check GPU
-        if self.device == "cuda":
-            if 'gpu_total_gb' in memory_info:
-                if memory_info['gpu_total_gb'] < ModelConfig.MIN_GPU_MEMORY:
-                    warnings.append(f"⚠️ Low GPU memory: {memory_info['gpu_total_gb']:.1f}GB (minimum {ModelConfig.MIN_GPU_MEMORY}GB recommended)")
-            else:
-                warnings.append("⚠️ Could not detect GPU memory")
-        
-        if warnings:
-            self.logger.warning("\n".join(warnings))
-        
-        return len(warnings) == 0
-    
-    def _download_model(self, repo_id: str, filename: str, repo_type: str = "model") -> Path:
-        """
-        Download model from Hugging Face Hub
-        
-        Args:
-            repo_id: Repository ID on Hugging Face
-            filename: Name of the file to download
-            repo_type: Type of repository
-            
-        Returns:
-            Path to downloaded file
-        """
-        try:
-            self.logger.info(f"Downloading {filename} from {repo_id}...")
-            
-            local_path = hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                repo_type=repo_type,
-                cache_dir=str(self.cache_dir),
-                resume_download=True
-            )
-            
-            self.logger.info(f"✓ Downloaded to: {local_path}")
-            return Path(local_path)
-            
-        except Exception as e:
-            error_msg = f"Failed to download {filename}: {str(e)}"
-            self.logger.error(error_msg)
-            self.load_errors.append(error_msg)
-            raise
-    
-    def _load_sam2(self) -> bool:
-        """
-        Load SAM2 model
-        
-        Returns:
-            Success status
-        """
-        try:
-            self.logger.info("Loading SAM2 model...")
-            
-            # Download checkpoint and config
-            checkpoint_path = self._download_model(
-                ModelConfig.SAM2_MODEL_ID,
-                ModelConfig.SAM2_CHECKPOINT
-            )
-            
-            config_path = self._download_model(
-                ModelConfig.SAM2_MODEL_ID,
-                ModelConfig.SAM2_CONFIG
-            )
-            
-            # Verify files exist
-            if not checkpoint_path.exists():
-                raise FileNotFoundError(f"SAM2 checkpoint not found: {checkpoint_path}")
-            if not config_path.exists():
-                raise FileNotFoundError(f"SAM2 config not found: {config_path}")
-            
-            # Build predictor
-            self.sam2_predictor = build_sam2_video_predictor(
-                str(config_path),
-                str(checkpoint_path),
-                device=self.device
-            )
-            
-            if self.sam2_predictor is None:
-                raise RuntimeError("SAM2 predictor initialization returned None")
-            
-            self.logger.info("✓ SAM2 model loaded successfully")
-            return True
-            
-        except Exception as e:
-            error_msg = f"Failed to load SAM2: {str(e)}"
-            self.logger.error(error_msg)
-            self.logger.debug(traceback.format_exc())
-            self.load_errors.append(error_msg)
-            return False
-    
-    def _load_matanyone(self) -> bool:
-        """
-        Load MatAnyone model
-        
-        Returns:
-            Success status
-        """
-        try:
-            self.logger.info("Loading MatAnyone model...")
-            
-            # Download checkpoint
-            checkpoint_path = self._download_model(
-                ModelConfig.MATANYONE_REPO,
-                ModelConfig.MATANYONE_CHECKPOINT
-            )
-            
-            if not checkpoint_path.exists():
-                raise FileNotFoundError(f"MatAnyone checkpoint not found: {checkpoint_path}")
-            
-            # Initialize model
-            model_init = ModelInitializer(device=self.device)
-            self.matanyone_model = model_init.setup_models(str(checkpoint_path))
-            
-            if self.matanyone_model is None:
-                raise RuntimeError("MatAnyone model initialization returned None")
-            
-            self.logger.info("✓ MatAnyone model loaded successfully")
-            return True
-            
-        except Exception as e:
-            error_msg = f"Failed to load MatAnyone: {str(e)}"
-            self.logger.error(error_msg)
-            self.logger.debug(traceback.format_exc())
-            self.load_errors.append(error_msg)
-            return False
-    
-    def load_models(self, force_reload: bool = False) -> bool:
-        """
-        Load all models
-        
-        Args:
-            force_reload: Force reload even if already loaded
-            
-        Returns:
-            Success status
-        """
-        if self.models_loaded and not force_reload:
-            self.logger.info("Models already loaded")
-            return True
-        
-        self.logger.info("Starting model loading process...")
-        self.logger.info(format_memory_status(get_memory_info()))
-        
-        # Clean memory before loading
-        clean_memory()
-        
-        # Reset status
-        self.models_loaded = False
-        self.load_errors = []
-        
-        # Load models
-        sam2_success = self._load_sam2()
-        matanyone_success = self._load_matanyone()
-        
-        # Update status
-        self.models_loaded = sam2_success and matanyone_success
-        
-        # Report results
-        if self.models_loaded:
-            self.logger.info("✅ All models loaded successfully")
-            self.logger.info(format_memory_status(get_memory_info()))
-        else:
-            self.logger.error("❌ Some models failed to load:")
-            for error in self.load_errors:
-                self.logger.error(f"  - {error}")
-        
-        return self.models_loaded
-    
-    def validate_models(self) -> Tuple[bool, Dict[str, Any]]:
-        """
-        Validate loaded models
-        
-        Returns:
-            Tuple of (success status, validation results)
-        """
-        results = {
-            'sam2': {'loaded': False, 'functional': False, 'error': None},
-            'matanyone': {'loaded': False, 'functional': False, 'error': None},
-            'memory': get_memory_info()
-        }
-        
-        # Check SAM2
-        if self.sam2_predictor is not None:
-            results['sam2']['loaded'] = True
-            try:
-                # Simple functionality test
-                test_frames = torch.randn(1, 3, 256, 256).to(self.device)
-                with torch.no_grad():
-                    # Just check if the model responds without error
-                    _ = self.sam2_predictor.model.image_encoder(test_frames)
-                results['sam2']['functional'] = True
-            except Exception as e:
-                results['sam2']['error'] = str(e)
-                self.logger.error(f"SAM2 validation failed: {e}")
-        
-        # Check MatAnyone
-        if self.matanyone_model is not None:
-            results['matanyone']['loaded'] = True
-            try:
-                # Simple functionality test
-                test_input = torch.randn(1, 4, 256, 256).to(self.device)
-                with torch.no_grad():
-                    # Just check if the model responds without error
-                    _ = self.matanyone_model(test_input)
-                results['matanyone']['functional'] = True
-            except Exception as e:
-                results['matanyone']['error'] = str(e)
-                self.logger.error(f"MatAnyone validation failed: {e}")
-        
-        # Overall success
-        success = (results['sam2']['functional'] and 
-                  results['matanyone']['functional'])
-        
-        return success, results
-    
-    def cleanup(self):
-        """Clean up models and free memory"""
-        self.logger.info("Cleaning up models...")
-        
-        # Delete model instances
-        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
-        
-        # Clear memory
-        clean_memory()
-        
-        self.models_loaded = False
-        self.logger.info("✓ Cleanup complete")
-    
-    def get_load_summary(self) -> str:
-        """Get summary of loaded models"""
-        lines = ["=== Model Load Summary ==="]
-        
-        if self.models_loaded:
-            lines.append("✅ Models loaded successfully")
-            lines.append(f"  - SAM2: {'Loaded' if self.sam2_predictor else 'Not loaded'}")
-            lines.append(f"  - MatAnyone: {'Loaded' if self.matanyone_model else 'Not loaded'}")
-        else:
-            lines.append("❌ Models not fully loaded")
-            if self.load_errors:
-                lines.append("Errors:")
-                for error in self.load_errors:
-                    lines.append(f"  - {error}")
-        
-        lines.append("")
-        lines.append(format_memory_status(get_memory_info()))
-        
-        return "\n".join(lines)
-    
-    def get_matanyone(self):
-        """Get MatAnyone model for backward compatibility"""
-        return self.matanyone_model
-    
-    def get_sam2(self):
-        """Get SAM2 predictor for backward compatibility"""
-        return self.sam2_predictor
-
-# ============================================================================
-# MAIN EXECUTION
-# ============================================================================
-
-def main():
-    """Test model loading"""
-    logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-    )
-    
-    print("Starting model loader test...")
-    print(f"Project root: {PROJECT_ROOT}")
-    print(f"Python path includes: {sys.path[:3]}")
-    
-    # Create loader
-    loader = ModelLoader()
-    
-    # Load models
-    success = loader.load_models()
-    
-    if success:
-        print("\n✅ Models loaded successfully!")
-        
-        # Validate models
-        valid, results = loader.validate_models()
-        
-        print("\nValidation Results:")
-        print(f"  SAM2: {results['sam2']}")
-        print(f"  MatAnyone: {results['matanyone']}")
-        
-        if valid:
-            print("\n✅ All models validated successfully!")
-        else:
-            print("\n⚠️ Some models failed validation")
-    else:
-        print("\n❌ Failed to load models")
-        print(loader.get_load_summary())
-    
-    # Cleanup
-    input("\nPress Enter to cleanup and exit...")
-    loader.cleanup()
-    print("Done!")
-
-if __name__ == "__main__":
-    main()