Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -26,6 +26,33 @@
 
 logger = logging.getLogger(__name__)
 
+# ============================================================================
+# LOADED MODEL DATA CONTAINER
+# ============================================================================
+class LoadedModel:
+    """
+    Tracks loaded model + metadata.
+    Useful for dashboards, export, analytics, etc.
+    """
+    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
+        self.load_time = load_time
+        self.device = device
+        self.framework = framework
+
+    def to_dict(self):
+        return {
+            "model_id": self.model_id,
+            "framework": self.framework,
+            "device": self.device,
+            "load_time": self.load_time,
+            "loaded": self.model is not None
+        }
+
+    def __repr__(self):
+        return f"LoadedModel(id={self.model_id}, loaded={self.model is not None}, device={self.device}, framework={self.framework}, load_time={self.load_time:.2f}s)"
+
 # ============================================================================
 # MODEL LOADER CLASS
 # ============================================================================
@@ -34,14 +61,13 @@ class ModelLoader:
     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()
 
-        self.sam2_predictor = None
-        self.matanyone_model = None  # This is usually InferenceCore
+        self.sam2_predictor = None  # LoadedModel instance or None
+        self.matanyone_model = None  # LoadedModel instance or None
 
         self.checkpoints_dir = "./checkpoints"
         os.makedirs(self.checkpoints_dir, exist_ok=True)
@@ -78,12 +104,13 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
             logger.info("Loading SAM2 predictor...")
             if progress_callback:
                 progress_callback(0.1, "Loading SAM2 predictor...")
-            self.sam2_predictor = self._load_sam2_predictor(progress_callback)
+            sam2_loaded = self._load_sam2_predictor(progress_callback)
 
-            if self.sam2_predictor is None:
+            if sam2_loaded is None:
                 logger.warning("SAM2 loading failed - will use fallback segmentation")
             else:
-                sam2_time = time.time() - start_time
+                self.sam2_predictor = sam2_loaded
+                sam2_time = self.sam2_predictor.load_time
                 self.loading_stats['sam2_load_time'] = sam2_time
                 logger.info(f"SAM2 loaded in {sam2_time:.2f}s")
 
@@ -93,19 +120,20 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
                 progress_callback(0.6, "Loading MatAnyOne model...")
             matanyone_start = time.time()
 
-            self.matanyone_model = self._load_matanyone_model(progress_callback)
+            matanyone_loaded = self._load_matanyone_model(progress_callback)
 
-            if self.matanyone_model is None:
+            if matanyone_loaded is None:
                 logger.warning("MatAnyOne loading failed - will use OpenCV refinement")
             else:
-                matanyone_time = time.time() - matanyone_start
+                self.matanyone_model = matanyone_loaded
+                matanyone_time = self.matanyone_model.load_time
                 self.loading_stats['matanyone_load_time'] = matanyone_time
                 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
+            self.loading_stats['models_loaded'] = bool(self.sam2_predictor or self.matanyone_model)
 
             if progress_callback:
                 if self.sam2_predictor or self.matanyone_model:
@@ -115,7 +143,7 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
 
             logger.info(f"Model loading completed in {total_time:.2f}s")
 
-            return self.sam2_predictor, self.matanyone_model
+            return (self.sam2_predictor, self.matanyone_model)
 
         except Exception as e:
             error_msg = f"Model loading failed: {str(e)}"
@@ -132,7 +160,7 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
     def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
         """
         Loads SAM2 using the official Hugging Face interface.
-        Returns: SAM2 predictor object or None
+        Returns: LoadedModel instance or None
         """
         model_size = "large"
         try:
@@ -161,11 +189,19 @@ def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
 
         try:
             from sam2.sam2_image_predictor import SAM2ImagePredictor
+            t0 = time.time()
             predictor = SAM2ImagePredictor.from_pretrained(model_id)
             if hasattr(predictor, 'model'):
                 predictor.model = predictor.model.to(self.device)
+            t1 = time.time()
             logger.info("SAM2 loaded successfully via official from_pretrained")
-            return predictor
+            return LoadedModel(
+                model=predictor,
+                model_id=model_id,
+                load_time=t1-t0,
+                device=str(self.device),
+                framework="sam2"
+            )
         except ImportError:
             logger.error("SAM2 module not found. Install with: pip install sam2")
             return None
@@ -179,34 +215,30 @@ def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
     def _load_matanyone_model(self, progress_callback: Optional[Callable] = 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.)
-        ---------------------------------------------
+        Returns: LoadedModel instance or None
         """
         try:
             if progress_callback:
                 progress_callback(0.7, "Loading MatAnyOne model...")
 
-            # --- HIGHLIGHT: SET ANY MatAnyOne SETTINGS HERE ---
             from matanyone import InferenceCore
-
-            # Example: To set chunk size or custom model repo, add kwargs here.
-            # See: https://huggingface.co/PeiqingYang/MatAnyone for config options
-
+            t0 = time.time()
             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
+                dtype=torch.float32,               # Or torch.float16 for fast, but only for GPUs with good fp16
                 # chunk_size=512,                  # Optional: for memory tuning on large videos
             )
-
             processor = InferenceCore(**matanyone_kwargs)
+            t1 = time.time()
             logger.info("MatAnyOne loaded successfully (InferenceCore)")
-            return processor
-
+            return LoadedModel(
+                model=processor,
+                model_id=matanyone_kwargs["repo_id"],
+                load_time=t1-t0,
+                device=str(self.device),
+                framework="matanyone"
+            )
         except ImportError:
             logger.error("MatAnyOne module not found. Install with: pip install matanyone")
             return None
@@ -245,9 +277,11 @@ def get_model_info(self) -> Dict[str, Any]:
             'loading_stats': self.loading_stats.copy()
         }
         if self.sam2_predictor is not None:
-            info['sam2_model_type'] = type(self.sam2_predictor).__name__
+            info['sam2_model_type'] = type(self.sam2_predictor.model).__name__
+            info['sam2_metadata'] = self.sam2_predictor.to_dict()
         if self.matanyone_model is not None:
-            info['matanyone_model_type'] = type(self.matanyone_model).__name__
+            info['matanyone_model_type'] = type(self.matanyone_model.model).__name__
+            info['matanyone_metadata'] = self.matanyone_model.to_dict()
         return info
 
     def get_load_summary(self) -> str:
@@ -258,11 +292,11 @@ def get_load_summary(self) -> str:
         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"
+            summary += f"✓ SAM2: {sam2_time:.1f}s (ID: {self.sam2_predictor.model_id})\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 (ID: {self.matanyone_model.model_id})\n"
         else:
             summary += f"✗ MatAnyOne: Failed (using OpenCV)\n"
         summary += f"Device: {self.device}"
@@ -278,7 +312,8 @@ def validate_models(self) -> bool:
         try:
             has_valid_model = False
             if self.sam2_predictor is not None:
-                if hasattr(self.sam2_predictor, 'set_image') or hasattr(self.sam2_predictor, 'predict'):
+                model = self.sam2_predictor.model
+                if hasattr(model, 'set_image') or hasattr(model, 'predict'):
                     has_valid_model = True
             if self.matanyone_model is not None:
                 has_valid_model = True
@@ -300,4 +335,3 @@ def models_ready(self) -> bool:
 # ============================================================================
 # END MODEL LOADER
 # ============================================================================
-