Update weight loading

Author: deruyter92

fmpose3d/fmpose3d.py

@@ -480,6 +480,9 @@ class _IngestedInput:
 # ---------------------------------------------------------------------------
 
 
+# FIXME @deruyter92: THIS IS TEMPORARY UNTIL WE DOWNLOAD THE WEIGHTS FROM HUGGINGFACE
+SKIP_WEIGHTS_VALIDATION = object() # sentinel value to indicate that the weights should not be validated
+
 class FMPose3DInference:
     """High-level, two-step inference API for FMPose3D.
 
@@ -534,7 +537,7 @@ def __init__(
         self,
         model_cfg: FMPose3DConfig | None = None,
         inference_cfg: InferenceConfig | None = None,
-        model_weights_path: str = "",
+        model_weights_path: str | Path | None = SKIP_WEIGHTS_VALIDATION,
         device: str | torch.device | None = None,
         *,
         estimator_2d: HRNetEstimator | SuperAnimalEstimator | None = None,
@@ -544,6 +547,9 @@ def __init__(
         self.inference_cfg = inference_cfg or InferenceConfig()
         self.model_weights_path = model_weights_path
 
+        # Validate model weights path (download if needed)
+        self._resolve_model_weights_path()
+
         # Skeleton configuration from the model config.
         self._joints_left: list[int] = list(self.model_cfg.joints_left)
         self._joints_right: list[int] = list(self.model_cfg.joints_right)
@@ -572,7 +578,7 @@ def __init__(
     @classmethod
     def for_animals(
         cls,
-        model_weights_path: str = "",
+        model_weights_path: str = SKIP_WEIGHTS_VALIDATION,
         *,
         device: str | torch.device | None = None,
         inference_cfg: InferenceConfig | None = None,
@@ -915,35 +921,46 @@ def _load_weights(self) -> None:
         state-dict keys and pull matching entries from the checkpoint so that
         extra keys in the checkpoint are silently ignored.
         """
-        if not self.model_weights_path:
-            raise ValueError(
-                "No model weights path provided. Pass 'model_weights_path' "
-                "to the FMPose3DInference constructor."
-            )
-        weights = Path(self.model_weights_path)
-        if not weights.exists():
-            raise ValueError(
-                f"Model weights file not found: {weights}. "
-                "Please provide a valid path to a .pth checkpoint file in the "
-                "FMPose3DInference constructor."
-            )
         if self._model_3d is None:
             raise ValueError("Model not initialised. Call setup_runtime() first.")
-        pre_dict = torch.load(
-            self.model_weights_path,
+        weights = self._resolve_model_weights_path()
+        state_dict = torch.load(
+            weights,
             weights_only=True,
             map_location=self.device,
         )
-        model_dict = self._model_3d.state_dict()
-        for name in model_dict:
-            if name in pre_dict:
-                model_dict[name] = pre_dict[name]
-        self._model_3d.load_state_dict(model_dict)
+        self._model_3d.load_state_dict(state_dict)
 
     # ------------------------------------------------------------------
     # Private helpers – input resolution
     # ------------------------------------------------------------------
 
+    def _resolve_model_weights_path(self) -> None:
+        # TODO @deruyter92: THIS IS TEMPORARY UNTIL WE DOWNLOAD THE WEIGHTS FROM HUGGINGFACE
+        if self.model_weights_path is SKIP_WEIGHTS_VALIDATION:
+            return SKIP_WEIGHTS_VALIDATION
+        
+        if not self.model_weights_path:
+            self._download_model_weights()
+        self.model_weights_path = Path(self.model_weights_path).resolve()
+        if not self.model_weights_path.exists():
+            raise ValueError(
+                f"Model weights file not found: {self.model_weights_path}. "
+                "Please provide a valid path to a .pth checkpoint file in the "
+                "FMPose3DInference constructor. Or leave it empty to download "
+                "the weights from huggingface."
+            )
+        return self.model_weights_path
+
+    def _download_model_weights(self) -> None:
+        """Download model weights from huggingface."""
+        # TODO @deruyter92: Implement download from huggingface
+        raise NotImplementedError(
+            "Downloading model weights from huggingface is not implemented yet."
+            "Please provide a valid path to a .pth checkpoint file in the "
+            "FMPose3DInference constructor."
+        )
+
     def _ingest_input(self, source: Source) -> _IngestedInput:
         """Normalise *source* into a ``(N, H, W, C)`` frames array.
 

tests/test_config.py

@@ -36,7 +36,7 @@
 class TestFMPose3DConfig:
     def test_defaults(self):
         cfg = FMPose3DConfig()
-        assert cfg.layers == 3
+        assert cfg.layers == 5
         assert cfg.channel == 512
         assert cfg.d_hid == 1024
         assert cfg.n_joints == 17

tests/test_fmpose3d.py

@@ -547,18 +547,18 @@ def test_corrupt_image_raises(self, api, tmp_path):
 
 class TestLoadWeightsErrors:
     def test_empty_path_raises(self):
-        api = FMPose3DInference(model_weights_path="", device="cpu")
-        api._model_3d = _ZeroVelocityModel()
-        with pytest.raises(ValueError, match="No model weights path"):
+        with pytest.raises((ValueError, NotImplementedError)):
+            api = FMPose3DInference(model_weights_path="", device="cpu")
+            api._model_3d = _ZeroVelocityModel()
             api._load_weights()
 
     def test_nonexistent_file_raises(self):
-        api = FMPose3DInference(
-            model_weights_path="/nonexistent/weights.pth",
-            device="cpu",
-        )
-        api._model_3d = _ZeroVelocityModel()
         with pytest.raises(ValueError, match="Model weights file not found"):
+            api = FMPose3DInference(
+            model_weights_path="/nonexistent/weights.pth",
+                device="cpu",
+            )
+            api._model_3d = _ZeroVelocityModel()
             api._load_weights()
 
     def test_model_not_initialized_raises(self, tmp_path):