Update src/generate_utils/generate.py

src/generate_utils/generate.py

@@ -26,8 +26,8 @@ class PressureGenerator:
     }
 
     def __init__(self, 
-                 ckpt_dir="generate_utils/lib/ckpt/pressurepose_20251222_180032",
-                 smpl_model_dir="E:/pyku/smpl_models",
+                 ckpt_dir="src/generate_utils/lib/ckpt/pressurepose_20251222_180032",
+                 smpl_model_dir="src/smpl_models",
                  device="cpu"):
         """
         初始化生成器：加载配置、权重和 SMPL 模型。
@@ -58,7 +58,13 @@ class PressureGenerator:
     def _load_config(self):
         config_path = os.path.join(self.ckpt_dir, 'config.yaml')
         if not os.path.exists(config_path):
-            raise FileNotFoundError(f"Config not found at {config_path}")
+            from huggingface_hub import snapshot_download
+            snapshot_download(
+                "yolozyk/PaGe",
+                local_dir="src/generate_utils/lib/ckpt/",
+                local_dir_use_symlinks=False,
+                ignore_patterns=["*.safetensors", ".gitattributes"],
+            )
         with open(config_path, 'r') as f:
             return yaml.safe_load(f)
 
@@ -67,7 +73,13 @@ class PressureGenerator:
         ckpt_path = os.path.join(self.ckpt_dir, 'ckpts', 'best_model.pth')
         
         if not os.path.exists(ckpt_path):
-            raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")
+            from huggingface_hub import snapshot_download
+            snapshot_download(
+                "yolozyk/PaGe",
+                local_dir="src/generate_utils/lib/ckpt/",
+                local_dir_use_symlinks=False,
+                ignore_patterns=["*.safetensors", ".gitattributes"],
+            )
 
         checkpoint = torch.load(ckpt_path, map_location=self.device)
         model.load_state_dict(checkpoint['model_state_dict'])