FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
ARG DEBIAN_FRONTEND=noninteractive

ENV TORCH_CUDA_ARCH_LIST "6.0;6.1;7.0;7.5;8.0"

RUN apt-get update -y \
    && apt-get install -y --no-install-recommends \
       build-essential \
       apt-utils \
       wget \
       ca-certificates \
       libglib2.0-0 \
       libsm6 \
       libxext6 \
       libxrender-dev \
       bmon \
       iotop \
       g++ \
       git \
       openjdk-11-jre-headless \
       libgl1-mesa-glx \
       libeigen3-dev \
       bzip2 \
       curl \
       libosmesa6-dev \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
    /bin/bash /tmp/miniconda.sh -b -p /opt/conda && \
    rm /tmp/miniconda.sh && \
    /opt/conda/bin/conda clean --all --yes && \
    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc

ENV PATH /opt/conda/bin:$PATH
RUN conda create -n maple python=3.10 -y && \
    echo "conda activate maple" >> ~/.bashrc

SHELL ["conda", "run", "-n", "maple", "/bin/bash", "-c"]

RUN pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
RUN pip install nuscenes-devkit av2
RUN pip install spconv-cu121
RUN pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.1+cu121.html
RUN pip install cupy-cuda12x==12.1.0

WORKDIR /maple
COPY requirements.txt /maple/requirements.txt
RUN pip install -r requirements.txt

COPY sam_hq2 /maple/sam_hq2
RUN cd sam_hq2 && pip install -e . && pip install --no-build-isolation -e grounding_dino && cd ..
COPY InvSR /maple/InvSR
RUN cd InvSR && pip install -e ".[torch]" && cd ..

WORKDIR /maple
CMD ["bash", "-c", "source /opt/conda/etc/profile.d/conda.sh && conda activate maple && exec bash"]