Update demo.py

demo.py

@@ -222,6 +222,8 @@ class BaseTrainer(object):
             
         return vq_models
     
+
+
     def _create_rvqvae_model(self, dim_pose: int, body_part: str) -> RVQVAE:
         """Create a single RVQVAE model with specified configuration."""
         args = self.args
@@ -230,12 +232,21 @@ class BaseTrainer(object):
             args.down_t, args.stride_t, args.width, args.depth,
             args.dilation_growth_rate, args.vq_act, args.vq_norm
         )
-        
-        # Load pretrained weights
-        checkpoint_path = getattr(args, f'vqvae_{body_part}_path')
-        state = torch.load(checkpoint_path, map_location='cpu')
-        model.load_state_dict(state['net'])
+
+        # Base directory = folder where demo.py lives
+        base_dir = Path(__file__).resolve().parent
+        checkpoint_path = base_dir / "ckpt" / f"net_300000_{body_part}.pth"
+
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(
+                f"RVQVAE checkpoint for '{body_part}' not found at '{checkpoint_path}'.\n"
+                f"CWD is {Path.cwd()}."
+            )
+
+        state = torch.load(str(checkpoint_path), map_location="cpu")
+        model.load_state_dict(state["net"])
         return model
+
       
     
     def inverse_selection(self, filtered_t, selection_array, n):