Add support for custom RF-DETR model upload via direct upload or HF Hub repository

README.md

@@ -51,15 +51,28 @@ This application is designed to run on Hugging Face Spaces. The following files
 
 ## Model Loading
 
-**RF-DETR Model:**
-- Upload your trained `rf-detr-medium.pth` file to the Space
-- The application will automatically find and load it
-
-**MedGemma Models:**
+**MedGemma Models (Automatic):**
 - Models download automatically from Hugging Face Hub on first use
 - No manual installation required
 - Choose between 4B (faster) or 27B (more accurate) models
 
+**RF-DETR Model (Your Custom Model):**
+You have two options for uploading your custom RF-DETR model:
+
+### Option 1: Direct Upload (Simple)
+1. Upload your `rf-detr-medium.pth` file directly to your Space
+2. The app will automatically find and use it
+
+### Option 2: Model Repository (Recommended)
+1. Create a separate Hugging Face model repository (e.g., `your-username/rf-detr-medical`)
+2. Upload your model files there
+3. Set the environment variable `RFDETR_HF_REPO` to your repository ID
+
+**To set the environment variable:**
+- Go to your Space settings
+- Add `RFDETR_HF_REPO` with your model repository ID (e.g., `your-username/rf-detr-medical`)
+- The app will download from your repository automatically
+
 ## Space Configuration
 
 For optimal performance, configure your Space settings:

app.py

@@ -90,8 +90,25 @@ memory_manager = MemoryManager()
 
 def find_checkpoint() -> Optional[str]:
     """Find RF-DETR checkpoint in various locations."""
+    # Check for HuggingFace model repository first (recommended)
+    import os
+    hf_model_id = os.environ.get("RFDETR_HF_REPO")
+    if hf_model_id:
+        try:
+            from huggingface_hub import hf_hub_download
+            print(f"Downloading RF-DETR from HuggingFace Hub: {hf_model_id}")
+            checkpoint_path = hf_hub_download(
+                repo_id=hf_model_id,
+                filename="rf-detr-medium.pth",
+                cache_dir="/tmp/hf_cache"
+            )
+            return checkpoint_path
+        except Exception as e:
+            print(f"Failed to download from HF Hub: {e}")
+
+    # Fallback to local files
     candidates = [
-        "rf-detr-medium.pth",  # Current directory
+        "rf-detr-medium.pth",  # Current directory (direct upload)
         "/tmp/results/checkpoint_best_total.pth",
         "/tmp/results/checkpoint_best_ema.pth",
         "/tmp/results/checkpoint_best_regular.pth",