Fix HF Spaces startup timeout by pre-caching models

Dockerfile

@@ -32,6 +32,16 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Increase Hugging Face Hub download timeout to prevent 504 errors
+ENV HF_HUB_DOWNLOAD_TIMEOUT=120
+
+# Create scripts directory and copy just the preload script first to leverage Docker cache
+RUN mkdir -p scripts
+COPY scripts/preload_models.py scripts/
+
+# Pre-download HF model weights into the image to prevent container startup timeouts
+RUN python scripts/preload_models.py
+
 # Copy the current directory contents into the container at /app
 COPY . .
 

scripts/preload_models.py

@@ -0,0 +1,43 @@
+import os
+import torch
+from transformers import (
+    DetrImageProcessor, 
+    DetrForObjectDetection,
+    ViTImageProcessor,
+    ViTForImageClassification,
+    CLIPProcessor,
+    CLIPModel,
+    AutoTokenizer,
+    AutoModel,
+    AutoModelForQuestionAnswering,
+    AutoModelForSeq2SeqLM,
+    BartForConditionalGeneration
+)
+
+# Set timeout
+os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"
+
+def preload():
+    print("🚀 Starting model pre-loading...")
+    
+    models = {
+        "detection": ("facebook/detr-resnet-50", DetrForObjectDetection, DetrImageProcessor),
+        "reid": ("google/vit-base-patch16-224", ViTForImageClassification, ViTImageProcessor),
+        "clip": ("openai/clip-vit-base-patch32", CLIPModel, CLIPProcessor),
+        "search": ("sentence-transformers/all-MiniLM-L6-v2", AutoModel, AutoTokenizer),
+        "qa": ("deepset/roberta-base-squad2", AutoModelForQuestionAnswering, AutoTokenizer),
+        "report": ("google/flan-t5-base", AutoModelForSeq2SeqLM, AutoTokenizer),
+        "summarizer": ("facebook/bart-large-cnn", BartForConditionalGeneration, AutoTokenizer),
+    }
+
+    for name, (model_id, model_cls, proc_cls) in models.items():
+        print(f"📦 Pre-loading {name}: {model_id}...")
+        try:
+            proc_cls.from_pretrained(model_id)
+            model_cls.from_pretrained(model_id)
+            print(f"✅ {name} loaded.")
+        except Exception as e:
+            print(f"❌ Failed to load {name}: {e}")
+
+if __name__ == "__main__":
+    preload()