Update app.py

app.py

@@ -28,11 +28,17 @@ class Config:
 
 class ModelManager:
     @staticmethod
-    def load_model(checkpoint_name: str):
-        model_path = hf_hub_download(
-            repo_id="shimu0215/seg", 
-            filename="sapiens_2b_render_people_epoch_25_torchscript.pt2",    
+    def load_model(checkpoint_name: str, type='depth'):
+        if checkpoint_name == 'seg':
+            model_path = hf_hub_download(
+            repo_id="shimu0215/seg",  # 你的模型仓库
+            filename="sapiens_1b_seg_foreground_epoch_8_torchscript.pt2",    # 你的模型文件
         )
+        else:
+            model_path = hf_hub_download(
+                repo_id="shimu0215/seg", 
+                filename="sapiens_2b_render_people_epoch_25_torchscript.pt2",    
+            )
         model = torch.jit.load(model_path)
         model.eval()
         model.to("cuda")
@@ -60,7 +66,7 @@ class ImageProcessor:
         depth_map = depth_output.squeeze().cpu().numpy()
 
         if seg_model_name != "no-bg-removal":
-            seg_model = ModelManager.load_model(Config.SEG_CHECKPOINTS[seg_model_name])
+            seg_model = ModelManager.load_model(Config.SEG_CHECKPOINTS[seg_model_name],type='seg')
             seg_output = ModelManager.run_model(seg_model, input_tensor, image.height, image.width)
             seg_mask = (seg_output.argmax(dim=1) > 0).float().cpu().numpy()[0]
             depth_map[seg_mask == 0] = np.nan