Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -2,6 +2,14 @@
 """
 Unified Model Loader
 Coordinates separate SAM2 and MatAnyone loaders for cleaner architecture
+
+Notes:
+- SAM2: exposes set_image(...) and predict(...)
+- MatAnyone: our loader returns a stateful callable adapter:
+    - callable(adapter) -> frame0: adapter(image_rgb01, mask01), frames>0: adapter(image_rgb01)
+    - optional: adapter.reset() to clear per-video memory
+  We therefore validate MatAnyone by checking "callable(...)" and/or presence of "reset",
+  not only ".step/.process".
 """
 
 from __future__ import annotations
@@ -27,7 +35,7 @@
 
 class LoadedModel:
     """Container for loaded model information"""
-    def __init__(self, model=None, model_id: str = "", load_time: float = 0.0, 
+    def __init__(self, model=None, model_id: str = "", load_time: float = 0.0,
                  device: str = "", framework: str = ""):
         self.model = model
         self.model_id = model_id
@@ -42,25 +50,26 @@ def to_dict(self) -> Dict[str, Any]:
             "device": self.device,
             "load_time": self.load_time,
             "loaded": self.model is not None,
+            "model_type": type(self.model).__name__ if self.model is not None else None,
         }
 
 
 class ModelLoader:
     """Main model loader that coordinates SAM2 and MatAnyone loaders"""
-    
+
     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()
-        
+
         # Initialize specialized loaders
         self.sam2_loader = SAM2Loader(device=str(self.device))
         self.matanyone_loader = MatAnyoneLoader(device=str(self.device))
-        
+
         # Model storage
         self.sam2_predictor: Optional[LoadedModel] = None
         self.matanyone_model: Optional[LoadedModel] = None
-        
+
         # Statistics
         self.loading_stats = {
             "sam2_load_time": 0.0,
@@ -69,7 +78,7 @@ def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
             "models_loaded": False,
             "loading_attempts": 0,
         }
-        
+
         logger.info(f"ModelLoader initialized for device: {self.device}")
 
     def load_all_models(
@@ -79,33 +88,29 @@ def load_all_models(
     ) -> Tuple[Optional[LoadedModel], Optional[LoadedModel]]:
         """
         Load all models using specialized loaders
-        
-        Args:
-            progress_callback: Optional callback for progress updates
-            cancel_event: Optional threading.Event for cancellation
-            
+
         Returns:
             Tuple of (sam2_model, 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...")
-                
+
             # Clean up any existing models
             self._cleanup_models()
-            
-            # Load SAM2
+
+            # -------------------- Load SAM2 -------------------- #
             if progress_callback:
                 progress_callback(0.1, "Loading SAM2 model...")
-            
+
             sam2_start = time.time()
             sam2_model = self.sam2_loader.load()
             sam2_time = time.time() - sam2_start
-            
+
             if sam2_model:
                 self.sam2_predictor = LoadedModel(
                     model=sam2_model,
@@ -118,21 +123,21 @@ def load_all_models(
                 logger.info(f"SAM2 loaded in {sam2_time:.2f}s")
             else:
                 logger.warning("SAM2 loading failed")
-                
-            # Check for cancellation
+
+            # Cancellation check
             if cancel_event and cancel_event.is_set():
                 if progress_callback:
                     progress_callback(1.0, "Model loading cancelled")
                 return self.sam2_predictor, None
-                
-            # Load MatAnyone
+
+            # ----------------- Load MatAnyone ------------------ #
             if progress_callback:
                 progress_callback(0.6, "Loading MatAnyone model...")
-                
+
             matanyone_start = time.time()
-            matanyone_model = self.matanyone_loader.load()
+            matanyone_model = self.matanyone_loader.load()  # returns stateful callable adapter or None
             matanyone_time = time.time() - matanyone_start
-            
+
             if matanyone_model:
                 self.matanyone_model = LoadedModel(
                     model=matanyone_model,
@@ -145,31 +150,30 @@ def load_all_models(
                 logger.info(f"MatAnyone loaded in {matanyone_time:.2f}s")
             else:
                 logger.warning("MatAnyone loading failed")
-                
-            # Update statistics
+
+            # ----------------- Finalize stats ------------------ #
             total_time = time.time() - start_time
             self.loading_stats["total_load_time"] = total_time
             self.loading_stats["models_loaded"] = bool(self.sam2_predictor or self.matanyone_model)
-            
-            # Final progress update
+
             if progress_callback:
                 if self.loading_stats["models_loaded"]:
                     progress_callback(1.0, "Models loaded successfully")
                 else:
                     progress_callback(1.0, "Model loading completed with failures")
-                    
+
             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(error_msg)
             self._cleanup_models()
             self.loading_stats["models_loaded"] = False
-            
+
             if progress_callback:
                 progress_callback(1.0, f"Error: {error_msg}")
-                
+
             return None, None
 
     def reload_models(
@@ -192,70 +196,105 @@ def get_sam2(self):
         return self.sam2_predictor.model if self.sam2_predictor else None
 
     def get_matanyone(self):
-        """Get MatAnyone processor model"""
+        """
+        Get MatAnyone processor model.
+
+        IMPORTANT: This returns the stateful callable adapter from MatAnyoneLoader:
+            - callable(image_rgb01[, mask01]) -> 2D alpha
+            - optional .reset() to clear memory per video
+        """
         return self.matanyone_model.model if self.matanyone_model else None
 
     def validate_models(self) -> bool:
         """Validate that loaded models have expected interfaces"""
         try:
             valid = False
-            
+
+            # Validate SAM2
             if self.sam2_predictor:
                 model = self.sam2_predictor.model
                 if hasattr(model, "set_image") and hasattr(model, "predict"):
                     valid = True
                     logger.info("SAM2 model validated")
-                    
+
+            # Validate MatAnyone (stateful adapter OR raw core)
             if self.matanyone_model:
                 model = self.matanyone_model.model
-                if hasattr(model, "step") or hasattr(model, "process"):
+                if callable(model):
+                    valid = True
+                    logger.info("MatAnyone adapter validated (callable)")
+                elif hasattr(model, "step") or hasattr(model, "process"):
+                    valid = True
+                    logger.info("MatAnyone core validated (.step/.process)")
+                elif hasattr(model, "reset"):
+                    # still accept an adapter exposing reset but not callable (unlikely)
                     valid = True
-                    logger.info("MatAnyone model validated")
-                    
+                    logger.info("MatAnyone object validated via reset()")
+                else:
+                    logger.warning("MatAnyone present but interface not recognized")
+
             return valid
-            
+
         except Exception as e:
             logger.error(f"Model validation failed: {e}")
             return False
 
     def get_model_info(self) -> Dict[str, Any]:
         """Get detailed information about loaded models"""
-        info = {
+        info: Dict[str, Any] = {
             "models_loaded": self.loading_stats["models_loaded"],
             "device": str(self.device),
             "loading_stats": self.loading_stats.copy(),
         }
-        
+
         # Add SAM2 info
         info["sam2"] = self.sam2_loader.get_info() if self.sam2_loader else {}
-        
-        # Add MatAnyone info
-        info["matanyone"] = self.matanyone_loader.get_info() if self.matanyone_loader else {}
-        
+
+        # Add MatAnyone info (augment with interface hints)
+        mat_info = self.matanyone_loader.get_info() if self.matanyone_loader else {}
+        try:
+            m = self.get_matanyone()
+            mat_info["callable"] = bool(callable(m))
+            mat_info["has_reset"] = bool(hasattr(m, "reset"))
+            mat_info["has_step"] = bool(hasattr(m, "step"))
+            mat_info["has_process"] = bool(hasattr(m, "process"))
+        except Exception:
+            pass
+        info["matanyone"] = mat_info
+
         return info
 
     def get_load_summary(self) -> str:
         """Get human-readable loading summary"""
         if not self.loading_stats["models_loaded"]:
             return "No models loaded"
-            
+
         lines = []
         lines.append(f"Models loaded in {self.loading_stats['total_load_time']:.1f}s")
-        
+
         if self.sam2_predictor:
             lines.append(f"✓ SAM2: {self.loading_stats['sam2_load_time']:.1f}s")
             lines.append(f"  Model: {self.sam2_predictor.model_id}")
         else:
             lines.append("✗ SAM2: Failed to load")
-            
+
         if self.matanyone_model:
-            lines.append(f"✓ MatAnyone: {self.loading_stats['matanyone_load_time']:.1f}s")
+            # Describe adapter/callable for clarity
+            iface = []
+            m = self.matanyone_model.model
+            if callable(m): iface.append("callable")
+            if hasattr(m, "reset"): iface.append("reset")
+            if hasattr(m, "step"): iface.append("step")
+            if hasattr(m, "process"): iface.append("process")
+            iface_str = f" ({', '.join(iface)})" if iface else ""
+
+            lines.append(f"✓ MatAnyone: {self.loading_stats['matanyone_load_time']:.1f}s{iface_str}")
             lines.append(f"  Model: {self.matanyone_model.model_id}")
         else:
             lines.append("✗ MatAnyone: Failed to load")
-            
+
         lines.append(f"Device: {self.device}")
-        
+
         return "\n".join(lines)
 
     def cleanup(self):
@@ -269,21 +308,27 @@ def _cleanup_models(self):
         if self.sam2_loader:
             self.sam2_loader.cleanup()
         if self.sam2_predictor:
-            del self.sam2_predictor
+            try:
+                del self.sam2_predictor
+            except Exception:
+                pass
             self.sam2_predictor = None
-            
+
         # Clean up MatAnyone
         if self.matanyone_loader:
             self.matanyone_loader.cleanup()
         if self.matanyone_model:
-            del self.matanyone_model
+            try:
+                del self.matanyone_model
+            except Exception:
+                pass
             self.matanyone_model = None
-            
+
         # Clear CUDA cache
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-            
+
         # Garbage collection
         gc.collect()
-        
-        logger.debug("Model cleanup completed")
+
+        logger.debug("Model cleanup completed")