fix(models): load trusted checkpoints with weights_only disabled

model_loader.py

@@ -196,6 +196,14 @@ def _extract_state_dict(checkpoint):
     return checkpoint
 
 
+def _load_checkpoint(path):
+    """Load trusted local checkpoints in a PyTorch-version-safe way."""
+    try:
+        return torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=device)
+
+
 def _looks_like_legacy_unet(state_dict):
     """Detect legacy UNet checkpoints by key namespace."""
     if not isinstance(state_dict, dict):
@@ -266,7 +274,7 @@ def load_models():
     # Load UNet
     try:
         if os.path.exists(UNET_PATH):
-            checkpoint = torch.load(UNET_PATH, map_location=device)
+            checkpoint = _load_checkpoint(UNET_PATH)
             unet_state = _extract_state_dict(checkpoint)
             loaded_unet_model = None
 
@@ -290,7 +298,7 @@ def load_models():
     # Load EfficientNet
     try:
         if os.path.exists(EFFICIENTNET_PATH):
-            checkpoint = torch.load(EFFICIENTNET_PATH, map_location=device)
+            checkpoint = _load_checkpoint(EFFICIENTNET_PATH)
             eff_state = _extract_state_dict(checkpoint)
             nested_classifier = isinstance(eff_state, dict) and 'classifier.1.1.weight' in eff_state
             loaded_efficientnet_model = create_efficientnet_model(num_classes=1, nested_classifier=nested_classifier)