Update app.py

app.py

@@ -10,6 +10,7 @@ from torchvision import transforms
 from PIL import Image
 import tempfile
 from classes_and_palettes import GOLIATH_PALETTE, GOLIATH_CLASSES
+from huggingface_hub import hf_hub_download
 
 class Config:
     # os.system(f"wget https://huggingface.co/shimu0215/seg/resolve/main/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2")
@@ -28,8 +29,16 @@ class Config:
 class ModelManager:
     @staticmethod
     def load_model(checkpoint_name: str):
-        checkpoint_path = os.path.join('shimu0215/seg', Config.CHECKPOINTS[checkpoint_name])
-        model = torch.jit.load(checkpoint_path)
+
+        
+
+        # 下载模型到本地（默认存储在 ~/.cache/huggingface/hub）
+        model_path = hf_hub_download(
+            repo_id="shimu0215/seg",  # 你的模型仓库
+            filename="sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2",    # 你的模型文件
+        )
+        # checkpoint_path = os.path.join('shimu0215/seg', Config.CHECKPOINTS[checkpoint_name])
+        model = torch.jit.load(model_path)
         model.eval()
         model.to("cuda")
         return model