Improve: Add better error handling and logging for model loading

model_server.py

@@ -94,32 +94,55 @@ class ProtectedModelServer:
             # Add segment-anything-2 to path (internally only)
             base_dir = Path(__file__).parent
             sam2_path = base_dir / "segment-anything-2"
+            
+            if not sam2_path.exists():
+                raise FileNotFoundError(f"SAM2 installation not found at {sam2_path}")
+            
             sys.path.insert(0, str(sam2_path))
             
-            from sam2.build_sam import build_sam2
-            from sam2.sam2_image_predictor import SAM2ImagePredictor
+            try:
+                from sam2.build_sam import build_sam2
+                from sam2.sam2_image_predictor import SAM2ImagePredictor
+            except ImportError as e:
+                raise ImportError("SAM2 not properly installed. Check build logs.") from e
             
             # Get paths internally - NEVER sent to client
             model_cfg = get_model_config_path()
             sam2_checkpoint = get_model_checkpoint_path()
-            fine_tuned_weights = get_finetuned_weights_path()
             
-            # Load model
+            # Load device
             device = "cuda" if torch.cuda.is_available() else "cpu"
+            print(f"Loading model on device: {device}")
             
+            # Load base SAM2 model
+            print(f"Loading SAM2 from {sam2_checkpoint}")
             self._model = build_sam2(model_cfg, sam2_checkpoint, device=device)
             self._predictor = SAM2ImagePredictor(self._model)
             
-            # Load fine-tuned weights
-            state_dict = torch.load(fine_tuned_weights, map_location=device)
-            self._predictor.model.load_state_dict(state_dict)
+            # Try to load fine-tuned weights if available
+            try:
+                fine_tuned_weights = get_finetuned_weights_path()
+                print(f"Loading fine-tuned weights from {fine_tuned_weights}")
+                state_dict = torch.load(fine_tuned_weights, map_location=device)
+                self._predictor.model.load_state_dict(state_dict)
+                print("Fine-tuned weights loaded successfully")
+            except FileNotFoundError:
+                print("Warning: Fine-tuned weights not found. Using base SAM2 model.")
+                print("To use fine-tuned model, upload VREyeSAM_uncertainity_best.torch to Space Files")
+            except Exception as e:
+                print(f"Warning: Could not load fine-tuned weights: {e}")
+                print("Continuing with base SAM2 model")
             
             # Model is now loaded - weights are NOT accessible to clients
             self._predictor.model.eval()
+            print("Model loaded successfully")
             
             return True
         except Exception as e:
-            raise RuntimeError(f"Model initialization failed") from e
+            print(f"Error loading model: {e}")
+            import traceback
+            traceback.print_exc()
+            raise RuntimeError(f"Model initialization failed: {str(e)}") from e
     
     def predict(self, image: np.ndarray, num_samples: int = 30) -> Tuple[np.ndarray, np.ndarray]:
         """