Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -1,74 +1,51 @@
+#!/usr/bin/env python3
 """
 Model Loading Module
-Handles loading and validation of SAM2 and MatAnyone AI models
+Handles loading and validation of SAM2 and MatAnyOne AI models
+(Modern version for BackgroundFX Pro – only edit this file for model loading logic)
 """
 
-# ============================================================================ #
+# ============================================================================
 # 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
+from typing import Optional, Dict, Any, Tuple, Callable
 
 import torch
-from omegaconf import DictConfig, OmegaConf
 
-# Import modular components - Updated paths for BackgroundFX Pro structure
-from core.exceptions import ModelLoadingError  # Updated import
-from utils.hardware.device_manager import DeviceManager  # Updated import
-from utils.system.memory_manager import MemoryManager  # Updated import
+# Modular dependencies (adapt as your structure changes)
+from core.exceptions import ModelLoadingError
+from utils.hardware.device_manager import DeviceManager
+from utils.system.memory_manager import MemoryManager
 
 logger = logging.getLogger(__name__)
 
-# ============================================================================ #
-# DATA CONTAINER CLASSES
-# ============================================================================ #
-
-class LoadedModel:
-    """Container for a loaded model with metadata"""
-    def __init__(self, model=None, model_id: str = "", load_time: float = 0.0):
-        self.model = model
-        self.model_id = model_id
-        self.load_time = load_time
-        self.device = None
-        self.framework = None
-        
-    def __repr__(self):
-        return f"LoadedModel(id={self.model_id}, loaded={self.model is not None})"
-
-# ============================================================================ #
-# MODEL LOADER CLASS - MAIN INTERFACE
-# ============================================================================ #
-
+# ============================================================================
+# MODEL LOADER CLASS
+# ============================================================================
 class ModelLoader:
     """
-    Simplified model loading for SAM2 and MatAnyone
-    Uses only the working loading strategies without redundant attempts
+    Loads and manages SAM2 and MatAnyOne models.
+    Tune all model-specific logic/settings here.
     """
-    
+
     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()
-        
-        # Model storage
+
         self.sam2_predictor = None
-        self.matanyone_model = None
-        self.matanyone_core = None
-        
-        # Configuration paths
+        self.matanyone_model = None  # This is usually InferenceCore
+
         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,
@@ -76,111 +53,90 @@ def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
             'models_loaded': False,
             'loading_attempts': 0
         }
-        
-        logger.info(f"ModelLoader initialized for device: {self.device}")
 
-# ============================================================================ #
-# MAIN MODEL LOADING ORCHESTRATION
-# ============================================================================ #
+        logger.info(f"ModelLoader initialized for device: {self.device}")
 
-    def load_all_models(self, progress_callback: Optional[callable] = None, cancel_event=None) -> Tuple[Any, Any]:
+    # ============================================================================
+    # MAIN LOADING FUNCTION (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)
+        Loads both SAM2 and MatAnyOne models with error handling.
+        Returns: (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
             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:
                 logger.warning("SAM2 loading failed - will use fallback segmentation")
             else:
                 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...")
+
+            # Load MatAnyOne
+            logger.info("Loading MatAnyOne model...")
             if progress_callback:
-                progress_callback(0.6, "Loading MatAnyone model...")
-                
+                progress_callback(0.6, "Loading MatAnyOne model...")
             matanyone_start = time.time()
-            
-            self.matanyone_model, self.matanyone_core = self._load_matanyone_model(progress_callback)
-            
+
+            self.matanyone_model = self._load_matanyone_model(progress_callback)
+
             if self.matanyone_model is None:
-                logger.warning("MatAnyone loading failed - will use OpenCV refinement")
+                logger.warning("MatAnyOne loading failed - will use OpenCV refinement")
             else:
                 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
+                logger.info(f"MatAnyOne loaded in {matanyone_time:.1f}s")
+
+            # Final status
             total_time = time.time() - start_time
             self.loading_stats['total_load_time'] = total_time
             self.loading_stats['models_loaded'] = True
-            
+
             if progress_callback:
                 if self.sam2_predictor or self.matanyone_model:
                     progress_callback(1.0, "Models loaded (with fallbacks available)")
                 else:
                     progress_callback(1.0, "Using fallback methods (models failed to load)")
-                
+
             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()}")
-            
-            # 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 - DIRECT OFFICIAL APPROACH ONLY
-# ============================================================================ #
-
-    def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
+    # ============================================================================
+    # SAM2 LOADING (OFFICIAL FROM_PRETRAINED)
+    # ============================================================================
+    def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
         """
-        Load SAM2 using only the official from_pretrained method that works
-        
-        Args:
-            progress_callback: Progress update callback
-            
-        Returns:
-            SAM2 predictor or None
+        Loads SAM2 using the official Hugging Face interface.
+        Returns: SAM2 predictor object or None
         """
-        # Determine model size based on device memory
-        model_size = "large"  # default
-        if hasattr(self.device_manager, 'get_device_memory_gb'):
-            try:
+        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"
@@ -189,139 +145,98 @@ def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
                 elif memory_gb < 12:
                     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}")
-        
+        except Exception as e:
+            logger.warning(f"Could not determine device memory: {e}")
+
         model_map = {
             "tiny": "facebook/sam2.1-hiera-tiny",
-            "small": "facebook/sam2.1-hiera-small", 
+            "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} model...")
-        
-        # Use ONLY the official SAM2 from_pretrained method that works
+
         try:
-            logger.info(f"Loading SAM2 from {model_id}...")
             from sam2.sam2_image_predictor import SAM2ImagePredictor
-            
-            # This is the method that successfully downloads and loads the model
             predictor = SAM2ImagePredictor.from_pretrained(model_id)
-            
-            # Move to correct device if needed
             if hasattr(predictor, 'model'):
                 predictor.model = predictor.model.to(self.device)
-            
             logger.info("SAM2 loaded successfully via official from_pretrained")
             return predictor
-            
-        except ImportError as e:
-            logger.error(f"SAM2 module not found. Install with: pip install sam2")
+        except ImportError:
+            logger.error("SAM2 module not found. Install with: pip install sam2")
             return None
-            
         except Exception as e:
             logger.error(f"SAM2 loading failed: {e}")
-            # Try downloading checkpoint manually as fallback
-            try:
-                logger.info("Attempting manual checkpoint download...")
-                import urllib.request
-                
-                checkpoint_url = f"https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2.1_hiera_{model_size}.pt"
-                checkpoint_path = os.path.join(self.checkpoints_dir, f"sam2.1_hiera_{model_size}.pt")
-                
-                if not os.path.exists(checkpoint_path):
-                    logger.info(f"Downloading checkpoint from {checkpoint_url}")
-                    urllib.request.urlretrieve(checkpoint_url, checkpoint_path)
-                
-                # Try loading with downloaded checkpoint
-                predictor = SAM2ImagePredictor.from_pretrained(model_id, checkpoint=checkpoint_path)
-                logger.info("SAM2 loaded successfully with manual checkpoint")
-                return predictor
-                
-            except Exception as fallback_error:
-                logger.error(f"Manual checkpoint fallback also failed: {fallback_error}")
-                return None
-
-# ============================================================================ #
-# MATANYONE MODEL LOADING
-# ============================================================================ #
-
-    def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
+            return None
+
+    # ============================================================================
+    # MATANYONE LOADING (OFFICIAL INFERENCECORE)
+    # ============================================================================
+    def _load_matanyone_model(self, progress_callback: Optional[Callable] = None):
         """
-        Load MatAnyone model - try official method only
-        
-        Args:
-            progress_callback: Progress update callback
-            
-        Returns:
-            Tuple[model, core] or (None, None)
+        Loads MatAnyOne using Hugging Face official 'matanyone' package.
+        Returns: InferenceCore object or None
+
+        ---------- MATANYONE TUNING SECTION ----------
+        To adjust MatAnyOne settings, change arguments to InferenceCore below!
+        (e.g., for precision, model variant, device, chunk size, etc.)
+        ---------------------------------------------
         """
         try:
-            logger.info("Loading MatAnyone from HuggingFace...")
             if progress_callback:
-                progress_callback(0.7, "Loading MatAnyone model...")
-            
+                progress_callback(0.7, "Loading MatAnyOne model...")
+
+            # --- HIGHLIGHT: SET ANY MatAnyOne SETTINGS HERE ---
             from matanyone import InferenceCore
-            
-            # Initialize with the official model repo
-            processor = InferenceCore("PeiqingYang/MatAnyone")
-            
-            logger.info("MatAnyone loaded successfully")
-            return processor, processor
-            
+
+            # Example: To set chunk size or custom model repo, add kwargs here.
+            # See: https://huggingface.co/PeiqingYang/MatAnyone for config options
+
+            matanyone_kwargs = dict(
+                repo_id="PeiqingYang/MatAnyone",   # You can change to any compatible Hugging Face repo
+                device=self.device,                # Device to load on ("cuda" or "cpu")
+                dtype=torch.float32,               # Change to torch.float16 for faster inference on good GPUs
+                # chunk_size=512,                  # Optional: for memory tuning on large videos
+            )
+
+            processor = InferenceCore(**matanyone_kwargs)
+            logger.info("MatAnyOne loaded successfully (InferenceCore)")
+            return processor
+
         except ImportError:
-            logger.error("MatAnyone module not found. Install with: pip install matanyone")
-            return None, None
-            
+            logger.error("MatAnyOne module not found. Install with: pip install matanyone")
+            return None
         except Exception as e:
-            logger.error(f"MatAnyone loading failed: {e}")
-            return None, None
-
-# ============================================================================ #
-# MODEL MANAGEMENT AND CLEANUP
-# ============================================================================ #
+            logger.error(f"MatAnyOne loading failed: {e}")
+            return None
 
+    # ============================================================================
+    # 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
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         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
-# ============================================================================ #
-
+    # ============================================================================
+    # MODEL INFO AND VALIDATION
+    # ============================================================================
     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,
@@ -329,118 +244,60 @@ def get_model_info(self) -> Dict[str, Any]:
             '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__
-                if hasattr(self.sam2_predictor, 'model'):
-                    info['sam2_has_model'] = True
-                if hasattr(self.sam2_predictor, 'predictor'):
-                    info['sam2_has_predictor'] = True
-            except:
-                info['sam2_model_type'] = "Unknown"
-        
+            info['sam2_model_type'] = type(self.sam2_predictor).__name__
         if self.matanyone_model is not None:
-            try:
-                info['matanyone_model_type'] = type(self.matanyone_model).__name__
-            except:
-                info['matanyone_model_type'] = "Unknown"
-        
+            info['matanyone_model_type'] = type(self.matanyone_model).__name__
         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 in {total_time:.1f}s\n"
-        
         if self.sam2_predictor:
             summary += f"✓ SAM2: {sam2_time:.1f}s\n"
         else:
             summary += f"✗ SAM2: Failed (using fallback)\n"
-            
         if self.matanyone_model:
-            summary += f"✓ MatAnyone: {matanyone_time:.1f}s\n"
+            summary += f"✓ MatAnyOne: {matanyone_time:.1f}s\n"
         else:
-            summary += f"✗ MatAnyone: Failed (using OpenCV)\n"
-            
+            summary += f"✗ MatAnyOne: Failed (using OpenCV)\n"
         summary += f"Device: {self.device}"
-        
         return summary
-    
+
     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
 
-# ============================================================================ #
-# MODEL VALIDATION AND TESTING
-# ============================================================================ #
-
     def validate_models(self) -> bool:
-        """
-        Validate that models are properly loaded and functional
-        
-        Returns:
-            bool: True if at least one model is valid
-        """
         try:
             has_valid_model = False
-            
-            # Check SAM2
             if self.sam2_predictor is not None:
-                # Check for required methods/attributes
                 if hasattr(self.sam2_predictor, 'set_image') or hasattr(self.sam2_predictor, 'predict'):
                     has_valid_model = True
-                    logger.info("SAM2 validation passed")
-                elif hasattr(self.sam2_predictor, 'model'):
-                    has_valid_model = True
-                    logger.info("SAM2 model found")
-            
-            # Check MatAnyone
             if self.matanyone_model is not None:
                 has_valid_model = True
-                logger.info("MatAnyone validation passed")
-            
             return has_valid_model
-            
         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)
-        """
+    def reload_models(self, progress_callback: Optional[Callable] = None) -> Tuple[Any, Any]:
         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 at least one model is loaded and ready"""
-        return self.sam2_predictor is not None or self.matanyone_model is not None
+        return self.sam2_predictor is not None or self.matanyone_model is not None
+
+# ============================================================================
+# END MODEL LOADER
+# ============================================================================
+