Update model loading logic and add support for multiple checkpoints

.gitattributes

@@ -1,3 +1,4 @@
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,2 +1 @@
-.DS_Store
-*.pt2
+.DS_Store

app.py

@@ -93,25 +93,27 @@ def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_i
 
 # ----------------- MODEL ----------------- #
 
-URL = "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.3b/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2?download=true"
 CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints")
-model_path = os.path.join(CHECKPOINTS_DIR, "sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2")
+CHECKPOINTS = {
+    "0.3B": "sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2",
+    "0.6B": "sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2",
+    "1B": "sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2",
+    "2B": "sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2",
+}
+
 
-if not os.path.exists(model_path):
-    os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
-    import requests
+def load_model(checkpoint_name: str):
+    checkpoint_path = os.path.join(CHECKPOINTS_DIR, CHECKPOINTS[checkpoint_name])
+    model = torch.jit.load(checkpoint_path)
+    model.eval()
+    model.to("cuda")
 
-    response = requests.get(URL)
-    with open(model_path, "wb") as file:
-        file.write(response.content)
 
-model = torch.jit.load(model_path)
-model.eval()
-model.to("cuda")
+MODELS = {name: load_model(name) for name in CHECKPOINTS.keys()}
 
 
 @torch.inference_mode()
-def run_model(input_tensor, height, width):
+def run_model(model, input_tensor, height, width):
     output = model(input_tensor)
     output = torch.nn.functional.interpolate(output, size=(height, width), mode="bilinear", align_corners=False)
     _, preds = torch.max(output, 1)
@@ -129,9 +131,10 @@ transform_fn = transforms.Compose(
 
 
 @spaces.GPU
-def segment(image: Image.Image) -> Image.Image:
+def segment(image: Image.Image, model_name: str) -> Image.Image:
     input_tensor = transform_fn(image).unsqueeze(0).to("cuda")
-    preds = run_model(input_tensor, height=image.height, width=image.width)
+    model = MODELS[model_name]
+    preds = run_model(model, input_tensor, height=image.height, width=image.width)
     mask = preds.squeeze(0).cpu().numpy()
     mask_image = Image.fromarray(mask.astype("uint8"))
     blended_image = visualize_mask_with_overlay(image, mask_image, LABELS_TO_IDS, alpha=0.5)
@@ -161,6 +164,11 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome(radius_size=sizes.radi
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil", format="png")
+            model_name = gr.Dropdown(
+                label="Model Version",
+                choices=list(CHECKPOINTS.keys()),
+                value="0.3B",
+            )
 
             example_model = gr.Examples(
                 inputs=input_image,
@@ -178,7 +186,7 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome(radius_size=sizes.radi
 
     run_button.click(
         fn=segment,
-        inputs=[input_image],
+        inputs=[input_image, model_name],
         outputs=[result_image],
     )
 

assets/checkpoints/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:735a9a8d63fe8f3f6a4ca3d787de07e69b1f9708ad550e09bb33c9854b7eafbc
+size 1358871599

assets/checkpoints/sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86aa2cb9d7310ba1cb1971026889f1d10d80ddf655d6028aea060aae94d82082
+size 2685144079

assets/checkpoints/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33bba30f3de8d9cfd44e4eaa4817b1bfdd98c188edfc87fa7cc031ba0f4edc17
+size 4716314057

assets/checkpoints/sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f32f841135794327a434b79fd25c6cca24a72e098e314baa430be65e13dd0332
+size 8706612665

banner.html

@@ -17,7 +17,7 @@
   font-family: 'Trebuchet MS', 'Lucida Sans Unicode', 'Lucida Grande',
     'Lucida Sans', Arial, sans-serif;
 ">
-      Sapiens 0.3B: Body-part Segmentation
+      Sapiens: Body-part Segmentation
     </h1>
 
 

requirements.txt

@@ -4,5 +4,4 @@ torch
 torchvision
 matplotlib
 pillow
-requests
 spaces