Update Dockerfile

Dockerfile

@@ -1,47 +1,69 @@
-# Use a slim Python image for a smaller footprint
-FROM python:3.11-slim
+# =========================================================
+# 1. BASE IMAGE (CUDA + PyTorch preinstalled)
+# =========================================================
+FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
 
-# Set environment variables
-# PYTHONDONTWRITEBYTECODE: Prevents Python from writing .pyc files
-# PYTHONUNBUFFERED: Ensures logs are sent straight to terminal without buffering
+# =========================================================
+# 2. ENV VARIABLES
+# =========================================================
 ENV PYTHONDONTWRITEBYTECODE=1 \
     PYTHONUNBUFFERED=1 \
     PORT=7860
 
-# Install system dependencies
-# Added libmupdf-dev if you decide to use advanced PyMuPDF features, 
-# though the pip package usually bundles what it needs.
-# Install system dependencies
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends \
-        git \
-        build-essential \
-        libgl1 \
-        libglib2.0-0 \
-        libsm6 \
-        libxrender1 \
-        libxext6 \
-    && apt-get clean \
+# =========================================================
+# 3. SYSTEM DEPENDENCIES
+# =========================================================
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    build-essential \
+    ninja-build \
+    cmake \
+    libglib2.0-0 \
+    libsm6 \
+    libxrender1 \
+    libxext6 \
     && rm -rf /var/lib/apt/lists/*
 
-# Set the working directory
+# =========================================================
+# 4. WORK DIRECTORY
+# =========================================================
 WORKDIR /app
 
-# Install Python dependencies first (leverages Docker layer caching)
+# =========================================================
+# 5. COPY REQUIREMENTS FIRST (Docker cache optimization)
+# =========================================================
 COPY requirements.txt .
-RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir -r requirements.txt
-# Copy the rest of the application code
 
+# Upgrade pip
+RUN pip install --no-cache-dir --upgrade pip
+
+# =========================================================
+# 6. INSTALL PYTHON DEPENDENCIES
+# =========================================================
+RUN pip install --no-cache-dir -r requirements.txt
+
+# =========================================================
+# 7. INSTALL FLASH ATTENTION (IMPORTANT)
+# =========================================================
+RUN pip install flash-attn==2.5.8 --no-build-isolation
+
+# =========================================================
+# 8. COPY APPLICATION CODE
+# =========================================================
 COPY . .
 
-# Create a non-root user for security (Best Practice)
+# =========================================================
+# 9. CREATE NON-ROOT USER (SECURITY)
+# =========================================================
 RUN useradd -m appuser && chown -R appuser /app
 USER appuser
 
-# Expose the designated port
+# =========================================================
+# 10. EXPOSE PORT
+# =========================================================
 EXPOSE 7860
 
-# Run the application
-# We use the list form of CMD for  better signal handling (CTRL+C)
+# =========================================================
+# 11. START APPLICATION
+# =========================================================
 CMD ["sh", "-c", "uvicorn mainapp:app --host 0.0.0.0 --port ${PORT}"]