download HF models during Image Build

demo/download_models.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+"""Pre-download HuggingFace models for Visual RAG Toolkit.
+
+This script downloads models during Docker build to cache them in the image,
+avoiding download delays during container startup.
+"""
+
+import os
+import sys
+
+os.environ.setdefault("HF_HOME", "/app/.cache/huggingface")
+os.environ.setdefault("TRANSFORMERS_CACHE", "/app/.cache/huggingface")
+
+MODELS_TO_DOWNLOAD = [
+    "vidore/colpali-v1.3",
+    "vidore/colSmol-500M",
+]
+
+def download_colpali_models():
+    """Download ColPali models and their processors."""
+    print("=" * 60)
+    print("Downloading ColPali models for Visual RAG Toolkit")
+    print("=" * 60)
+    
+    try:
+        from colpali_engine.models import ColPali, ColPaliProcessor
+    except ImportError:
+        print("[WARN] colpali-engine not installed, trying transformers directly")
+        from transformers import AutoModel, AutoProcessor
+        
+        for model_name in MODELS_TO_DOWNLOAD:
+            print(f"\n[INFO] Downloading model: {model_name}")
+            try:
+                AutoModel.from_pretrained(model_name, trust_remote_code=True)
+                AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+                print(f"[OK] Downloaded: {model_name}")
+            except Exception as e:
+                print(f"[WARN] Could not download {model_name}: {e}")
+        return
+    
+    for model_name in MODELS_TO_DOWNLOAD:
+        print(f"\n[INFO] Downloading model: {model_name}")
+        try:
+            if "colsmol" in model_name.lower():
+                from colpali_engine.models import ColQwen2, ColQwen2Processor
+                ColQwen2.from_pretrained(model_name, trust_remote_code=True)
+                ColQwen2Processor.from_pretrained(model_name, trust_remote_code=True)
+            else:
+                ColPali.from_pretrained(model_name, trust_remote_code=True)
+                ColPaliProcessor.from_pretrained(model_name, trust_remote_code=True)
+            print(f"[OK] Downloaded: {model_name}")
+        except Exception as e:
+            print(f"[WARN] Could not download {model_name} with colpali-engine: {e}")
+            try:
+                from transformers import AutoModel, AutoProcessor
+                AutoModel.from_pretrained(model_name, trust_remote_code=True)
+                AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+                print(f"[OK] Downloaded via transformers: {model_name}")
+            except Exception as e2:
+                print(f"[ERROR] Failed to download {model_name}: {e2}")
+
+
+def main():
+    print(f"[INFO] HF_HOME: {os.environ.get('HF_HOME', 'not set')}")
+    print(f"[INFO] Cache dir: {os.environ.get('TRANSFORMERS_CACHE', 'not set')}")
+    
+    download_colpali_models()
+    
+    print("\n" + "=" * 60)
+    print("Model download complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()