Fix: Simplify Dockerfile and add runtime weight download to model_server

Dockerfile

@@ -38,30 +38,8 @@ COPY app.py model_server.py ./
 RUN mkdir -p .streamlit
 COPY .streamlit/config.toml .streamlit/
 
-# Download VREyeSAM fine-tuned weights from Hugging Face
-# This assumes your weights are available on HF Model Hub
-RUN if [ ! -z "$HF_TOKEN" ]; then \
-    python -c "
-from huggingface_hub import hf_hub_download
-import os
-
-hf_token = os.getenv('HF_TOKEN')
-if hf_token:
-    # Download VREyeSAM weights from your HF repo
-    # Replace 'your-username/vreyesam' with your actual repo
-    try:
-        checkpoint = hf_hub_download(
-            repo_id='devnagaich/VREyeSAM',
-            filename='VREyeSAM_uncertainity_best.torch',
-            token=hf_token,
-            cache_dir='segment-anything-2/checkpoints'
-        )
-        print(f'Downloaded VREyeSAM weights to: {checkpoint}')
-    except Exception as e:
-        print(f'Warning: Could not download VREyeSAM weights: {e}')
-        print('App will attempt to load from cache at runtime')
-" || true; \
-    fi
+# Note: VREyeSAM fine-tuned weights will be downloaded at runtime by model_server.py
+# using the HF_TOKEN from HF Spaces Secrets
 
 # Expose Streamlit port
 EXPOSE 7860

model_server.py

@@ -17,22 +17,54 @@ from typing import Tuple, Optional
 def get_model_checkpoint_path():
     """Get checkpoint path secretly, never expose to client"""
     base_dir = Path(__file__).parent
-    checkpoint = base_dir / "segment-anything-2" / "checkpoints" / "sam2_hiera_small.pt"
+    checkpoint = base_dir / "segment-anything-2" / "checkpoints" / "sam2.1_hiera_small.pt"
     if not checkpoint.exists():
         raise FileNotFoundError(f"Model checkpoint not found")
     return str(checkpoint)
 
 def get_finetuned_weights_path():
-    """Get fine-tuned weights path secretly, never expose to client"""
+    """Get fine-tuned weights path secretly, never expose to client
+    
+    Attempts to download from Hugging Face if local copy doesn't exist
+    and HF_TOKEN is available.
+    """
     base_dir = Path(__file__).parent
-    weights = base_dir / "segment-anything-2" / "checkpoints" / "VREyeSAM_uncertainity_best.torch"
-    if not weights.exists():
-        raise FileNotFoundError(f"Fine-tuned weights not found")
-    return str(weights)
+    checkpoint_dir = base_dir / "segment-anything-2" / "checkpoints"
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    weights = checkpoint_dir / "VREyeSAM_uncertainity_best.torch"
+    
+    # If weights already exist locally, return path
+    if weights.exists():
+        return str(weights)
+    
+    # Try to download from Hugging Face using HF_TOKEN
+    hf_token = os.getenv('HF_TOKEN', '')
+    if hf_token:
+        try:
+            from huggingface_hub import hf_hub_download
+            print("Downloading VREyeSAM weights from Hugging Face...")
+            
+            checkpoint_path = hf_hub_download(
+                repo_id='devnagaich/VREyeSAM',
+                filename='VREyeSAM_uncertainity_best.torch',
+                token=hf_token,
+                cache_dir=str(checkpoint_dir)
+            )
+            print(f"Successfully downloaded VREyeSAM weights")
+            return checkpoint_path
+        except Exception as e:
+            print(f"Warning: Could not download VREyeSAM weights: {e}")
+    
+    # If download fails or no token, return path anyway (may exist from upload)
+    if weights.exists():
+        return str(weights)
+    
+    # Last resort - raise error
+    raise FileNotFoundError(f"VREyeSAM weights not found and could not download")
 
 def get_model_config_path():
     """Get model config path secretly, never expose to client"""
-    return "configs/sam2/sam2_hiera_s.yaml"
+    return "configs/sam2.1/sam2.1_hiera_s.yaml"
 
 
 class ProtectedModelServer: