ARG TOKEN
FROM  pytorch/pytorch:2.2.0-cuda11.8-cudnn8-devel as first_build
ARG TOKEN
ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6+PTX" \
    TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
    CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
    FORCE_CUDA="1" \
    PATH=/usr/local/cuda/bin:$PATH
ENV LANG C.UTF-8
RUN sed -i "s@http://.*archive.ubuntu.com@http://mirrors.bfsu.edu.cn@g" /etc/apt/sources.list \
    && sed -i "s@http://.*security.ubuntu.com@http://mirrors.bfsu.edu.cn@g" /etc/apt/sources.list \
    # && rm /etc/apt/sources.list.d/cuda.list \
    # && rm /etc/apt/sources.list.d/nvidia-ml.list \
    # && apt-key del 7fa2af80 \
    # && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub \
    # && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub \
    && pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
RUN apt-get update \
    && apt-get install -y \
    ffmpeg \
    libsm6 \
    libxext6 \
    git \
    ninja-build \
    libglib2.0-0 \
    libsm6 \
    libxrender-dev \
    libxext6 \
    libgl1 \
    libglib2.0-0 \
    libturbojpeg \
    curl \
    vim \
    unzip \
    zip \
    rar \
    unrar \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/* \
    && echo -e "set termencoding=utf-8\nset encoding=utf8\nset fileencodings=utf8,ucs-bom,gbk,cp936,gb2312,gb18030" > /root/.vimrc \
    && LAZYGIT_VERSION=$(curl -s "https://api.github.com/repos/jesseduffield/lazygit/releases/latest" | grep -Po '"tag_name": "v\K[^"]*') \
    && curl -Lo lazygit.tar.gz "https://github.com/jesseduffield/lazygit/releases/latest/download/lazygit_${LAZYGIT_VERSION}_Linux_x86_64.tar.gz" \
    && tar xf lazygit.tar.gz lazygit \
    && rm lazygit.tar.gz \
    && install lazygit /usr/local/bin
RUN cd /root \
    && conda clean --all \
    && mkdir /root/mount/ \
    && pip install --no-cache-dir opencv-python-headless \
    && pip install timm==0.5.4 \
    && git clone https://${TOKEN}@git.chiebot.com:10000/HuQiong/mae-cnn.git MAE_CNN\
    && cd /root/MAE_CNN/pretrain \
    && pip install -r requirements.txt\
    && cd /root