make GPU compatible

Dockerfile

@@ -1,7 +1,9 @@
-FROM python:3.11-slim
+# Use HuggingFace's pre-built GPU image (includes CUDA, PyTorch, Transformers)
+FROM huggingface/transformers-pytorch-gpu:latest@sha256:4c7317881a534b22e18add49c925096fa902651fb0571c69f3cad58af3ea2c0f
 
 WORKDIR /app
 
+# Install system dependencies
 RUN apt-get update && apt-get install -y \
     build-essential \
     curl \
@@ -9,36 +11,32 @@ RUN apt-get update && apt-get install -y \
     poppler-utils \
     && rm -rf /var/lib/apt/lists/*
 
-# Set HuggingFace cache directories BEFORE downloading models
+# Verify Python version
+RUN python3 -V
+
+# Copy requirements first (for better Docker layer caching)
+COPY requirements.txt ./
+
+# Install Python dependencies
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Pre-download Hugging Face models during build
+# IMPORTANT: Set cache directory BEFORE downloading to ensure models are cached
 ENV HF_HOME=/app/.cache/huggingface
-ENV TRANSFORMERS_CACHE=/app/.cache/huggingface
-ENV HF_DATASETS_CACHE=/app/.cache/huggingface
 ENV HF_HUB_CACHE=/app/.cache/huggingface
+ENV HF_DATASETS_CACHE=/app/.cache/huggingface
+ENV TRANSFORMERS_CACHE=/app/.cache/huggingface
 
-# Install core dependencies first (for better Docker layer caching)
-RUN pip3 install --no-cache-dir \
-    torch>=2.0.0 \
-    numpy>=1.21.0 \
-    Pillow>=9.0.0 \
-    tqdm>=4.60.0 \
-    pyyaml>=6.0 \
-    python-dotenv>=0.19.0 \
-    colpali-engine>=0.3.0 \
-    transformers>=4.35.0 \
-    qdrant-client>=1.7.0 \
-    streamlit>=1.25.0 \
-    httpx>=0.24.0 \
-    pandas \
-    altair \
-    datasets
-
-# Pre-download HuggingFace models during build
 COPY download_models.py ./
 RUN mkdir -p /app/.cache/huggingface && \
     chmod -R 755 /app/.cache && \
-    python download_models.py
+    chmod -R 755 /app/.cache/huggingface && \
+    python3 download_models.py
+
+# Install colpali-engine (after model download to ensure deps are ready)
+RUN pip3 install colpali-engine einops~=0.8.1
 
-# Copy application code
+# Copy all application files
 COPY visual_rag/ ./visual_rag/
 COPY benchmarks/ ./benchmarks/
 COPY demo/ ./demo/
@@ -51,6 +49,7 @@ RUN pip3 install -e .
 # Setup Streamlit configuration
 RUN mkdir -p /app/.streamlit && \
     mkdir -p /app/.cache/streamlit && \
+    mkdir -p /app/.cache/huggingface && \
     mkdir -p /app/results && \
     printf '[server]\nport = 7860\nheadless = true\nenableCORS = false\nenableXsrfProtection = false\nmaxUploadSize = 500\n\n[browser]\ngatherUsageStats = false\n' > /app/.streamlit/config.toml && \
     chmod -R 777 /app/.streamlit && \
@@ -59,15 +58,14 @@ RUN mkdir -p /app/.streamlit && \
     chmod -R 777 /app/results
 
 ENV STREAMLIT_CONFIG_HOME=/app/.streamlit
-ENV STREAMLIT_USER_BASE_PATH=/app/.cache/streamlit
 ENV STREAMLIT_BROWSER_GATHER_USAGE_STATS=false
+ENV STREAMLIT_USER_BASE_PATH=/app/.cache/streamlit
 
-# HuggingFace Spaces uses port 7860
+# Expose port (HF Spaces uses 7860)
 EXPOSE 7860
 
-# Health check for Streamlit
-HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
-    CMD curl --fail http://localhost:7860/_stcore/health || exit 1
+# Verify GPU availability (will show False during build, True at runtime)
+RUN python3 -c "import torch; print('CUDA available:', torch.cuda.is_available())" || true
 
 # Run Streamlit app
 ENTRYPOINT ["streamlit", "run", "demo_app.py", "--server.port=7860", "--server.address=0.0.0.0", "--server.headless", "true"]