feat: patch SAM to load on CPU

src/snowleopard_reid/pipeline/stages/segmentation.py

@@ -70,7 +70,20 @@ def load_sam_predictor(
     device_str = get_device(device=device, verbose=True)
 
     logger.info(f"Loading SAM HQ model: {model_type}")
-    sam = sam_model_registry[model_type](checkpoint=str(checkpoint_path))
+
+    # Patch torch.load to handle CPU-only environments
+    # SAM HQ's registry doesn't accept map_location, so we patch temporarily
+    import torch
+
+    original_load = torch.load
+    torch.load = lambda f, *args, **kwargs: original_load(
+        f, *args, map_location=device_str, **kwargs
+    )
+    try:
+        sam = sam_model_registry[model_type](checkpoint=str(checkpoint_path))
+    finally:
+        torch.load = original_load
+
     sam.to(device=device_str)
 
     predictor = SamPredictor(sam)