Added hf model

app.py

@@ -1,5 +1,5 @@
 """
-HuggingFace Space for MedSAM Inference
+HuggingFace Space for SAM / MedSAM Inference
 API-compatible with Dense-Captioning-Toolkit backend
 
 Deploy this to: https://huggingface.co/spaces/YOUR_USERNAME/medsam-inference
@@ -12,30 +12,49 @@ import io
 import json
 import base64
 
-# Import MedSAM components
+from huggingface_hub import hf_hub_download
+
+# Import SAM components
 from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
 
 # Initialize model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
-# Model configuration - using MedSAM (vit_b) for both interactive and automatic segmentation
-MODEL_CHECKPOINT = "medsam_vit_b.pth"
-MODEL_TYPE = "vit_b"
+# -----------------------------------------------------------------------------
+# Model configuration
+# -----------------------------------------------------------------------------
+# We CANNOT store the large SAM checkpoint directly in the Space repo due to
+# size limits, so we fetch it from a separate model repo on HuggingFace Hub
+# (e.g. `Aniketg6/dense-captioning-models`) and let Spaces cache it.
+#
+# You said you've uploaded your SAM model there; by default we assume the
+# filename is `sam_vit_h_4b8939.pth`. If it's different, just change
+# MODEL_FILENAME below.
+MODEL_REPO_ID = "Aniketg6/dense-captioning-models"
+MODEL_FILENAME = "sam_vit_h_4b8939.pth"  # change if your filename is different
+MODEL_TYPE = "vit_h"  # using SAM ViT-H (general-purpose SAM)
+
+print(f"Downloading SAM checkpoint `{MODEL_FILENAME}` from repo `{MODEL_REPO_ID}`...")
+MODEL_CHECKPOINT = hf_hub_download(
+    repo_id=MODEL_REPO_ID,
+    filename=MODEL_FILENAME,
+)
+print(f"✓ Checkpoint downloaded to: {MODEL_CHECKPOINT}")
 
-print("Loading MedSAM model...")
+print("Loading SAM model...")
 
 # Monkey-patch torch.load to use CPU mapping when needed
 original_torch_load = torch.load
 def patched_torch_load(f, *args, **kwargs):
-    if 'map_location' not in kwargs and device == 'cpu':
-        kwargs['map_location'] = 'cpu'
+    if "map_location" not in kwargs and device == "cpu":
+        kwargs["map_location"] = "cpu"
     return original_torch_load(f, *args, **kwargs)
 
 torch.load = patched_torch_load
 
-# Load MedSAM model (vit_b) - used for both interactive and automatic segmentation
-print(f"Loading MedSAM model ({MODEL_TYPE})...")
+# Load SAM model (vit_h) - used for both interactive and automatic segmentation
+print(f"Loading SAM model ({MODEL_TYPE}) from downloaded checkpoint...")
 
 try:
     torch.load = patched_torch_load
@@ -66,7 +85,7 @@ mask_generator = SamAutomaticMaskGenerator(
     min_mask_region_area=0         # Minimum mask area (lowered from 100 to allow small masks)
 )
 print("✓ SamAutomaticMaskGenerator initialized for automatic segmentation")
-print("✓ MedSAM model loaded successfully!")
+print("✓ SAM model loaded successfully from HuggingFace Hub!")
 
 
 # =============================================================================
@@ -465,12 +484,12 @@ def check_auto_mask_status():
     """
     Check if automatic mask generation is available
     """
-    return json.dumps({
-        'available': mask_generator is not None,
-        'model': 'medsam_vit_b' if mask_generator else None,
-        'model_type': MODEL_TYPE,
-        'device': str(device)
-    })
+            return json.dumps({
+                'available': mask_generator is not None,
+                'model': MODEL_FILENAME if mask_generator else None,
+                'model_type': MODEL_TYPE,
+                'device': str(device)
+            })
 
 
 # =============================================================================