FROM nvidia/cuda:12.1.1-devel-ubuntu20.04
ENV NVIDIA_VISIBLE_DEVICES ${NVIDIA_VISIBLE_DEVICES:-all}
ENV NVIDIA_DRIVER_CAPABILITIES ${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}graphics

ARG PYTHON_VERSION=3.10

# Install os-level packages
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    bash-completion \
    build-essential \
    ca-certificates \
    cmake \
    curl \
    git \
    htop \
    libegl1 \
    libxext6 \
    libjpeg-dev \
    libpng-dev  \
    rsync \
    tmux \
    unzip \
    vim \
    wget \
    xvfb \
    && rm -rf /var/lib/apt/lists/*

# Install (mini) conda
RUN curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
    chmod +x ~/miniconda.sh && \
    ~/miniconda.sh -b -p /opt/conda && \
    rm ~/miniconda.sh && \
    /opt/conda/bin/conda init && \
    /opt/conda/bin/conda install -y python="$PYTHON_VERSION" && \
    /opt/conda/bin/conda clean -ya

ENV PATH /opt/conda/bin:$PATH
SHELL ["/bin/bash", "-c"]

RUN pip install \
    numpy==1.26.4 \
    scipy \
    ninja \
    torch==2.1.2 \
    torchvision==0.16.2 \
    h5py \
    matplotlib \
    "trimesh>=4.2.0" \
    "pyglet<2" \
    "accelerate>=0.28.0" \
    wandb \
    timm \
    datasets \
    hydra-core \
    && pip cache purge

RUN FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;9.0" pip install "git+https://github.com/Jiayuan-Gu/torkit3d.git@235ecf60497271136f5552cb45bb7cf75ab1cb09" && pip cache purge

# Install apex
RUN git clone --single-branch https://github.com/NVIDIA/apex && \
    cd apex && git checkout 810ffae374a2b9cb4b5c5e28eaeca7d7998fca0c && \
    pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ && pip cache purge && \
    cd .. && rm -rf apex

RUN useradd -m -u 1000 user

WORKDIR /app

RUN pip install git+https://github.com/zyc00/Point-SAM.git && pip cache purge
RUN pip install flask flask_cors && pip cache purge

COPY --chown=user . /app

RUN wget https://yuchen-service.nrp-nautilus.io/yuchen_fast/pointcloud-sam/pretrained/ours/mixture_10k/model-2.safetensors

CMD [ "python3", "app.py", "--host=0.0.0.0", "--port=7860"]