Update model_loader.py

model_loader.py

@@ -3,6 +3,10 @@
 Handles loading and validation of SAM2 and MatAnyone AI models
 """
 
+# ============================================================================ #
+# IMPORTS AND DEPENDENCIES
+# ============================================================================ #
+
 import os
 import gc
 import time
@@ -13,7 +17,6 @@
 from pathlib import Path
 
 import torch
-import hydra
 import gradio as gr
 from omegaconf import DictConfig, OmegaConf
 
@@ -24,9 +27,14 @@
 
 logger = logging.getLogger(__name__)
 
+# ============================================================================ #
+# 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):
@@ -40,7 +48,6 @@ def __init__(self, device_mgr: device_manager.DeviceManager, memory_mgr: memory_
         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)
         
@@ -55,6 +62,10 @@ def __init__(self, device_mgr: device_manager.DeviceManager, memory_mgr: memory_
         
         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"""
@@ -75,7 +86,11 @@ def patched_get_config(self):
             
         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
@@ -152,83 +167,69 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
                 progress_callback(1.0, f"Error: {error_msg}")
             
             return None, None
-    
+
+# ============================================================================ #
+# SAM2 MODEL LOADING - AUTOMATIC CONFIG DETECTION
+# ============================================================================ #
+
     def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
         """
-        Load SAM2 predictor with multiple fallback strategies
+        Load SAM2 predictor with automatic config detection - no manual config files needed
+        Uses build_sam2_video_predictor for automatic configuration based on checkpoint filename
         
         Args:
             progress_callback: Progress update callback
             
         Returns:
-            SAM2ImagePredictor or None
+            SAM2VideoPredictor or None
         """
-        if not os.path.isdir(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"""
+        def try_load_sam2_auto(repo_id: str, filename: str, model_name: str):
+            """Attempt to load SAM2 with automatic config detection"""
             try:
-                checkpoint_path = os.path.join(self.checkpoints_dir, checkpoint_name)
+                checkpoint_path = os.path.join(self.checkpoints_dir, filename)
                 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...")
+                    logger.info(f"Downloading {filename} from Hugging Face Hub...")
                     if progress_callback:
-                        progress_callback(0.2, f"Downloading {checkpoint_name}...")
+                        progress_callback(0.2, f"Downloading {filename}...")
                     
                     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,
+                            repo_id=repo_id,
+                            filename=filename,
                             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}")
+                        logger.warning(f"Failed to download {filename}: {download_error}")
                         return None
                 
-                # 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", "")
                 if progress_callback:
-                    progress_callback(0.4, f"Building {config_name}...")
+                    progress_callback(0.4, f"Building SAM2 {model_name}...")
                 
-                from sam2.build_sam import build_sam2
-                from sam2.sam2_image_predictor import SAM2ImagePredictor
+                # Use automatic config detection - NO manual config needed!
+                from sam2.build_sam import build_sam2_video_predictor
                 
-                sam2_model = build_sam2(config_name, checkpoint_path)
-                sam2_model.to(self.device)
-                predictor = SAM2ImagePredictor(sam2_model)
+                predictor = build_sam2_video_predictor(checkpoint_path, device=self.device)
                 
-                logger.info(f"SAM2 {config_name} loaded successfully on {self.device}")
+                logger.info(f"SAM2 {model_name} loaded successfully on {self.device}")
                 return predictor
                 
             except Exception as e:
-                error_msg = f"Failed to load SAM2 {config_name_with_yaml}: {e}"
+                error_msg = f"Failed to load SAM2 {model_name}: {e}"
                 logger.warning(error_msg)
                 return None
         
-        # Try different SAM2 model sizes based on device capabilities
+        # Try different SAM2 models with automatic config detection
         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")
+            ("facebook/sam2-hiera-large", "sam2_hiera_large.pt", "hiera_large"),
+            ("facebook/sam2-hiera-base-plus", "sam2_hiera_base_plus.pt", "hiera_base_plus"),
+            ("facebook/sam2-hiera-small", "sam2_hiera_small.pt", "hiera_small"),
+            ("facebook/sam2-hiera-tiny", "sam2_hiera_tiny.pt", "hiera_tiny")
         ]
         
         # Prioritize model size based on device memory
@@ -242,14 +243,18 @@ def try_load_sam2(config_name_with_yaml: str, checkpoint_name: str):
             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)
+        for repo_id, filename, model_name in model_attempts:
+            predictor = try_load_sam2_auto(repo_id, filename, model_name)
             if predictor is not None:
                 return predictor
         
         logger.error("All SAM2 model loading attempts 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
@@ -284,7 +289,11 @@ def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
         
         logger.error("All MatAnyone loading strategies failed")
         return None, None
-    
+
+# ============================================================================ #
+# MATANYONE LOADING STRATEGIES
+# ============================================================================ #
+
     def _load_matanyone_strategy_1(self):
         """MatAnyone loading strategy 1: Direct model import"""
         from matanyone.model.matanyone import MatAnyOne
@@ -351,7 +360,11 @@ def _load_matanyone_strategy_4(self):
         model = load_model_from_hub(model_path, device=self.device)
         
         return model, model  # Return same object for both
-    
+
+# ============================================================================ #
+# MODEL MANAGEMENT AND CLEANUP
+# ============================================================================ #
+
     def _cleanup_models(self):
         """Clean up loaded models and free memory"""
         if self.sam2_predictor is not None:
@@ -372,6 +385,15 @@ def _cleanup_models(self):
         
         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
@@ -420,7 +442,11 @@ def get_load_summary(self) -> str:
         summary += f"Device: {self.device}"
         
         return summary
-    
+
+# ============================================================================ #
+# MODEL VALIDATION AND TESTING
+# ============================================================================ #
+
     def validate_models(self) -> bool:
         """
         Validate that models are properly loaded and functional
@@ -444,7 +470,11 @@ def validate_models(self) -> bool:
         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)
@@ -461,11 +491,6 @@ def reload_models(self, progress_callback: Optional[callable] = None) -> Tuple[A
         
         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:
         """Check if all models are loaded and ready"""