Enable CUDA 12.6 runtime, default flash-attn on HF, add huggingface_hub/PyMuPDF, make flash-attn configurable

Dockerfile

@@ -1,4 +1,10 @@
-FROM python:3.11-slim
+FROM pytorch/pytorch:2.7.0-cuda12.6-cudnn9-runtime
+
+# Build args to optionally enable flash-attn installation and override wheel URL
+# Enable by default for Hugging Face Spaces GPU builds; override locally with
+#   --build-arg INSTALL_FLASH_ATTN=false
+ARG INSTALL_FLASH_ATTN=true
+ARG FLASH_ATTN_WHEEL_URL=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.8/flash_attn-2.7.4.post1+cu126torch2.7-cp310-cp310-linux_x86_64.whl
 
 # Install system dependencies as root
 RUN apt-get update && apt-get install -y \
@@ -35,13 +41,21 @@ ENV HOME=/home/user \
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app
 
-# Try and run pip command after setting the user with `USER user` to avoid permission issues with Python
+# Upgrade pip in the CUDA-enabled base image
 RUN pip install --no-cache-dir --upgrade pip
 
 # Copy requirements and install Python dependencies
 COPY --chown=user requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Optionally install flash-attn wheel (requires Python/torch/CUDA compatibility)
+RUN if [ "$INSTALL_FLASH_ATTN" = "true" ]; then \
+      echo "Installing flash-attn from $FLASH_ATTN_WHEEL_URL" && \
+      pip install --no-cache-dir $FLASH_ATTN_WHEEL_URL; \
+    else \
+      echo "Skipping flash-attn installation"; \
+    fi
+
 # Copy application code
 COPY --chown=user . .
 

requirements.txt

@@ -5,5 +5,5 @@ pydantic>=2.2.0,<3.0.0
 opencv-python>=4.9.0.80
 numpy>=1.26.0
 pillow>=10.3.0
-torch>=2.2.0
-torchvision>=0.17.0
+huggingface_hub
+PyMuPDF