Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -19,7 +19,6 @@
 from pathlib import Path
 
 import torch
-import gradio as gr
 from omegaconf import DictConfig, OmegaConf
 
 # Import modular components - Updated paths for BackgroundFX Pro structure
@@ -45,179 +44,14 @@ def __init__(self, model=None, model_id: str = "", load_time: float = 0.0):
     def __repr__(self):
         return f"LoadedModel(id={self.model_id}, loaded={self.model is not None})"
 
-# ============================================================================ #
-# 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
+    Simplified model loading for SAM2 and MatAnyone
+    Uses only the working loading strategies without redundant attempts
     """
     
     def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
@@ -244,31 +78,6 @@ def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
         }
         
         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
@@ -296,7 +105,7 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             # Clear any existing models
             self._cleanup_models()
             
-            # Load SAM2 first (typically faster)
+            # Load SAM2 first
             logger.info("Loading SAM2 predictor...")
             if progress_callback:
                 progress_callback(0.1, "Loading SAM2 predictor...")
@@ -304,11 +113,11 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             self.sam2_predictor = self._load_sam2_predictor(progress_callback)
             
             if self.sam2_predictor is None:
-                raise 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")
+                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...")
@@ -320,11 +129,11 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             self.matanyone_model, self.matanyone_core = self._load_matanyone_model(progress_callback)
             
             if self.matanyone_model is None:
-                raise 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")
+                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
             total_time = time.time() - start_time
@@ -332,9 +141,12 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             self.loading_stats['models_loaded'] = True
             
             if progress_callback:
-                progress_callback(1.0, "Models loaded successfully!")
+                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"All models loaded successfully in {total_time:.2f}s")
+            logger.info(f"Model loading completed in {total_time:.2f}s")
             
             return self.sam2_predictor, self.matanyone_model
             
@@ -352,29 +164,20 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             return None, None
 
 # ============================================================================ #
-# SAM2 MODEL LOADING - HUGGINGFACE TRANSFORMERS APPROACH
+# SAM2 MODEL LOADING - DIRECT OFFICIAL APPROACH ONLY
 # ============================================================================ #
 
     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
+        Load SAM2 using only the official from_pretrained method that works
         
         Args:
             progress_callback: Progress update callback
             
         Returns:
-            SAM2 model or None
+            SAM2 predictor 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
+        # Determine model size based on device memory
         model_size = "large"  # default
         if hasattr(self.device_manager, 'get_device_memory_gb'):
             try:
@@ -382,12 +185,13 @@ def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
                 if memory_gb < 4:
                     model_size = "tiny"
                 elif memory_gb < 8:
+                    model_size = "small"
+                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}")
         
-        # Step 3: Try multiple HuggingFace approaches
         model_map = {
             "tiny": "facebook/sam2.1-hiera-tiny",
             "small": "facebook/sam2.1-hiera-small", 
@@ -398,86 +202,57 @@ def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
         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}")
+            progress_callback(0.3, f"Loading SAM2 {model_size} model...")
         
-        # Method 2: Direct Transformers classes
+        # Use ONLY the official SAM2 from_pretrained method that works
         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...")
+            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 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 ImportError as e:
+            logger.error(f"SAM2 module not found. Install with: pip install sam2")
+            return None
             
         except Exception as e:
-            logger.warning(f"Fallback approach failed: {e}")
-        
-        logger.error("All SAM2 loading methods failed")
-        return None
+            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 - MULTIPLE STRATEGIES
+# MATANYONE MODEL LOADING
 # ============================================================================ #
 
     def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
         """
-        Load MatAnyone model with multiple import strategies
+        Load MatAnyone model - try official method only
         
         Args:
             progress_callback: Progress update callback
@@ -485,84 +260,26 @@ def _load_matanyone_model(self, progress_callback: Optional[callable] = None):
         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
+            logger.info("Loading MatAnyone from HuggingFace...")
+            if progress_callback:
+                progress_callback(0.7, "Loading MatAnyone model...")
+            
+            from matanyone import InferenceCore
+            
+            # Initialize with the official model repo
+            processor = InferenceCore("PeiqingYang/MatAnyone")
+            
+            logger.info("MatAnyone loaded successfully")
+            return processor, processor
+            
         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
+            logger.error("MatAnyone module not found. Install with: pip install matanyone")
+            return None, None
+            
+        except Exception as e:
+            logger.error(f"MatAnyone loading failed: {e}")
+            return None, None
 
 # ============================================================================ #
 # MODEL MANAGEMENT AND CLEANUP
@@ -583,7 +300,8 @@ def _cleanup_models(self):
             self.matanyone_core = None
         
         # Clear GPU cache
-        self.memory_manager.cleanup_aggressive()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         gc.collect()
         
         logger.debug("Model cleanup completed")
@@ -615,6 +333,10 @@ def get_model_info(self) -> Dict[str, Any]:
         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"
         
@@ -639,9 +361,18 @@ def get_load_summary(self) -> str:
         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"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"
+        else:
+            summary += f"✗ MatAnyone: Failed (using OpenCV)\n"
+            
         summary += f"Device: {self.device}"
         
         return summary
@@ -663,20 +394,27 @@ def validate_models(self) -> bool:
         Validate that models are properly loaded and functional
         
         Returns:
-            bool: True if models are valid
+            bool: True if at least one model is valid
         """
         try:
-            # Basic validation
-            if not self.loading_stats['models_loaded']:
-                return False
+            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")
             
-            if self.sam2_predictor is None or self.matanyone_model is None:
-                return False
+            # Check MatAnyone
+            if self.matanyone_model is not None:
+                has_valid_model = True
+                logger.info("MatAnyone validation passed")
             
-            # Try basic model operations
-            # This could include running a small test inference
-            logger.info("Model validation passed")
-            return True
+            return has_valid_model
             
         except Exception as e:
             logger.error(f"Model validation failed: {e}")
@@ -704,9 +442,5 @@ def reload_models(self, progress_callback: Optional[callable] = None) -> Tuple[A
     
     @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
-        )
+        """Check if at least one model is loaded and ready"""
+        return self.sam2_predictor is not None or self.matanyone_model is not None