FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive

# Install system dependencies needed for compiling extensions and running OpenCV
RUN apt-get update && apt-get install -y \
    git \
    python3.10 \
    python3.10-dev \
    python3-pip \
    ninja-build \
    build-essential \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Set 'python' command to use python3.10
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1

# Set up the non-root user required by Hugging Face Spaces
RUN useradd -m -u 1000 user
USER user
ENV PATH="/home/user/.local/bin:$PATH"

WORKDIR /home/user/app

# Install PyTorch with CUDA 12.1
RUN pip install --no-cache-dir torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu121

# Install standard requirements
COPY --chown=user requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy all the application files (Python codebase, submodules, outputs, assets)
COPY --chown=user . .

# Compile and install the custom CUDA extensions
ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0"
ENV MAX_JOBS=2
RUN pip install --no-cache-dir --no-build-isolation ./submodules/diff-gaussian-rasterization ./submodules/simple-knn git+https://github.com/NVlabs/nvdiffrast

# Set Gradio environmental variables
ENV GRADIO_SERVER_NAME="0.0.0.0"
ENV GRADIO_SERVER_PORT=7860

EXPOSE 7860

CMD ["python", "app.py"]