Fix: Load model with CPU mapping for HuggingFace Space

app.py

@@ -22,11 +22,25 @@ MODEL_CHECKPOINT = "medsam_vit_b.pth"
 MODEL_TYPE = "vit_b"
 
 print("Loading MedSAM model...")
-sam = sam_model_registry[MODEL_TYPE](checkpoint=MODEL_CHECKPOINT)
-sam.to(device=device)
-sam.eval()
-predictor = SamPredictor(sam)
-print("✓ MedSAM model loaded successfully!")
+
+# Monkey-patch torch.load to use CPU mapping when needed
+original_torch_load = torch.load
+def patched_torch_load(f, *args, **kwargs):
+    if 'map_location' not in kwargs and device == 'cpu':
+        kwargs['map_location'] = 'cpu'
+    return original_torch_load(f, *args, **kwargs)
+
+torch.load = patched_torch_load
+
+try:
+    sam = sam_model_registry[MODEL_TYPE](checkpoint=MODEL_CHECKPOINT)
+    sam.to(device=device)
+    sam.eval()
+    predictor = SamPredictor(sam)
+    print("✓ MedSAM model loaded successfully!")
+finally:
+    # Restore original torch.load
+    torch.load = original_torch_load
 
 
 def segment_with_points(image, points_json):