Update models/loaders/model_loader.py

models/loaders/model_loader.py

@@ -5,9 +5,6 @@
 (Modern version for BackgroundFX Pro – only edit this file for model loading logic)
 """
 
-# ============================================================================
-# IMPORTS AND DEPENDENCIES
-# ============================================================================
 import os
 import gc
 import sys
@@ -19,7 +16,6 @@
 
 import torch
 
-# 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
@@ -30,10 +26,6 @@
 # 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
@@ -57,10 +49,6 @@ def __repr__(self):
 # MODEL LOADER CLASS
 # ============================================================================
 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
@@ -86,10 +74,6 @@ def __init__(self, device_mgr: DeviceManager, memory_mgr: MemoryManager):
     # MAIN LOADING FUNCTION (ORCHESTRATION)
     # ============================================================================
     def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_event=None) -> Tuple[Any, Any]:
-        """
-        Loads both SAM2 and MatAnyOne models with error handling.
-        Returns: (sam2_predictor, matanyone_model)
-        """
         start_time = time.time()
         self.loading_stats['loading_attempts'] += 1
 
@@ -100,6 +84,9 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
 
             self._cleanup_models()
 
+            # --- DIAG: Log device and model selection step
+            logger.info(f"Device for models: {self.device}")
+
             # Load SAM2 first
             logger.info("Loading SAM2 predictor...")
             if progress_callback:
@@ -158,10 +145,6 @@ def load_all_models(self, progress_callback: Optional[Callable] = None, cancel_e
     # SAM2 LOADING (OFFICIAL FROM_PRETRAINED)
     # ============================================================================
     def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
-        """
-        Loads SAM2 using the official Hugging Face interface.
-        Returns: LoadedModel instance or None
-        """
         model_size = "large"
         try:
             if hasattr(self.device_manager, 'get_device_memory_gb'):
@@ -183,6 +166,7 @@ def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
             "large": "facebook/sam2.1-hiera-large"
         }
         model_id = model_map.get(model_size, model_map["large"])
+        logger.info(f"[DIAG] About to load SAM2 model_id: {model_id} on device {self.device}")
 
         if progress_callback:
             progress_callback(0.3, f"Loading SAM2 {model_size} model...")
@@ -191,6 +175,8 @@ def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
             from sam2.sam2_image_predictor import SAM2ImagePredictor
             t0 = time.time()
             predictor = SAM2ImagePredictor.from_pretrained(model_id)
+            logger.info(f"[DIAG] SAM2 predictor instance type: {type(predictor)}")
+            # If this fails, it's likely a missing model or bad download
             if hasattr(predictor, 'model'):
                 predictor.model = predictor.model.to(self.device)
             t1 = time.time()
@@ -202,21 +188,22 @@ def _load_sam2_predictor(self, progress_callback: Optional[Callable] = None):
                 device=str(self.device),
                 framework="sam2"
             )
+        except IndexError as e:
+            logger.error(f"SAM2 IndexError: {e}. (Did the model download fail? Wrong model_id?)")
+            logger.error(traceback.format_exc())
+            return None
         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}")
+            logger.error(traceback.format_exc())
             return None
 
     # ============================================================================
     # MATANYONE LOADING (OFFICIAL INFERENCECORE)
     # ============================================================================
     def _load_matanyone_model(self, progress_callback: Optional[Callable] = None):
-        """
-        Loads MatAnyOne using Hugging Face official 'matanyone' package.
-        Returns: LoadedModel instance or None
-        """
         try:
             if progress_callback:
                 progress_callback(0.7, "Loading MatAnyOne model...")
@@ -224,12 +211,14 @@ def _load_matanyone_model(self, progress_callback: Optional[Callable] = None):
             from matanyone import InferenceCore
             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,               # Or torch.float16 for fast, but only for GPUs with good fp16
-                # chunk_size=512,                  # Optional: for memory tuning on large videos
+                repo_id="PeiqingYang/MatAnyone",
+                device=self.device,
+                dtype=torch.float32,
+                # chunk_size=512,
             )
+            logger.info(f"[DIAG] About to load MatAnyOne from repo: {matanyone_kwargs['repo_id']} on device {self.device}")
             processor = InferenceCore(**matanyone_kwargs)
+            logger.info(f"[DIAG] MatAnyOne processor type: {type(processor)}")
             t1 = time.time()
             logger.info("MatAnyOne loaded successfully (InferenceCore)")
             return LoadedModel(
@@ -239,11 +228,16 @@ def _load_matanyone_model(self, progress_callback: Optional[Callable] = None):
                 device=str(self.device),
                 framework="matanyone"
             )
+        except IndexError as e:
+            logger.error(f"MatAnyOne IndexError: {e}. (Did the model download fail? Wrong repo_id?)")
+            logger.error(traceback.format_exc())
+            return None
         except ImportError:
             logger.error("MatAnyOne module not found. Install with: pip install matanyone")
             return None
         except Exception as e:
             logger.error(f"MatAnyOne loading failed: {e}")
+            logger.error(traceback.format_exc())
             return None
 
     # ============================================================================
@@ -335,3 +329,4 @@ def models_ready(self) -> bool:
 # ============================================================================
 # END MODEL LOADER
 # ============================================================================
+