ARG BASE_IMAGE=pytorch/pytorch:2.7.0-cuda12.6-cudnn9-runtime
FROM ${BASE_IMAGE}

# 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-cp311-cp311-linux_x86_64.whl

# Persist caches and model storage in Spaces, and enable fast transfers
ENV HF_HUB_ENABLE_HF_TRANSFER=1 \
    HUGGINGFACE_HUB_CACHE=/data/.cache/huggingface \
    HF_HOME=/data/.cache/huggingface \
    DOTS_OCR_LOCAL_DIR=/data/models/dots-ocr \
    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb=512 \
    OMP_NUM_THREADS=1

# Install system dependencies as root
RUN apt-get update && apt-get install -y \
    libgl1-mesa-dri \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender1 \
    libgomp1 \
    libgtk-3-0 \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    libv4l-dev \
    libxvidcore-dev \
    libx264-dev \
    libjpeg-dev \
    libpng-dev \
    libtiff-dev \
    python3-dev \
    && rm -rf /var/lib/apt/lists/*

# Set up a new user named "user" with user ID 1000
RUN useradd -m -u 1000 user

# Create persistent data directories and grant write access to the user
RUN mkdir -p /data/.cache/huggingface /data/models/dots-ocr && \
    chown -R 1000:1000 /data

# Switch to the "user" user
USER user

# Set home to the user's home directory
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

# Set the working directory to the user's home directory
WORKDIR $HOME/app

# 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

# Copy pyproject.toml for package installation
COPY --chown=user pyproject.toml .

# Optionally install flash-attn wheel (requires Python/torch/CUDA compatibility)
# Will auto-skip if the wheel's Python tag does not match this image's Python.
RUN if [ "$INSTALL_FLASH_ATTN" = "true" ]; then \
      PYTAG=$(python -c "import sys; print(f'cp{sys.version_info.major}{sys.version_info.minor}')"); \
      echo "Detected Python tag: $PYTAG"; \
      if echo "$FLASH_ATTN_WHEEL_URL" | grep -q "$PYTAG"; then \
        echo "Installing flash-attn from $FLASH_ATTN_WHEEL_URL" && \
        pip install --no-cache-dir "$FLASH_ATTN_WHEEL_URL"; \
      else \
        echo "flash-attn wheel tag mismatch for $PYTAG. Skipping flash-attn install."; \
      fi; \
    else \
      echo "Skipping flash-attn installation"; \
    fi

# Copy source code
COPY --chown=user src/ ./src/
COPY --chown=user main.py .

# Install the package in development mode
RUN pip install --no-cache-dir -e .

# Expose port
EXPOSE 7860

# Run the application
CMD ["python", "main.py"]