ARG BASE_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-cuda11.3.0-py37-torch1.11.0-tf1.15.5-base
FROM $BASE_IMAGE
RUN apt-get update && \
    apt-get install -y libsox-dev unzip  zip iputils-ping telnet sudo && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

ARG CUDA_VERSION=cu121
# install  jupyter plugin
RUN mkdir -p /root/.local/share/jupyter/labextensions/ && \
    cp -r  /tmp/resources/jupyter_plugins/*  /root/.local/share/jupyter/labextensions/

COPY docker/scripts/modelscope_env_init.sh /usr/local/bin/ms_env_init.sh
# python3.8 pip install git+https://github.com/jin-s13/xtcocoapi.git@v1.13
# pip install git+https://github.com/gatagat/lap.git@v0.4.0
RUN pip install --no-cache-dir numpy 'cython<=0.29.36' funtextprocessing kwsbp==0.0.6 safetensors typeguard==2.13.3 scikit-learn librosa==0.9.2 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html

RUN pip install --no-cache-dir adaseq text2sql_lgesql==1.3.0 \
         git+https://github.com/jin-s13/xtcocoapi.git@v1.14 \
         git+https://github.com/gatagat/lap.git@v0.4.0 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force --no-deps

RUN mv /opt/conda/compiler_compat/ld /opt/conda/compiler_compat/ldbk && \
         pip install --no-cache-dir mpi4py paint_ldm \
         mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 \
         ipykernel fasttext fairseq deepspeed -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html

ARG USE_GPU


RUN if [ "$USE_GPU" = "True" ] ; then \
        CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0" pip install --no-cache-dir  'git+https://github.com/facebookresearch/detectron2.git'; \
    else \
        echo 'cpu unsupport detectron2'; \
    fi

# torchmetrics==0.11.4 for ofa
# tinycudann for cuda12.1.0 pytorch 2.1.2
RUN if [ "$USE_GPU" = "True" ] ; then \
    pip install --no-cache-dir torchsde jupyterlab torchmetrics==0.11.4 tiktoken transformers_stream_generator bitsandbytes basicsr optimum && \
    pip install --no-cache-dir auto-gptq --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu121/ && \
    pip install --no-cache-dir -U 'xformers<0.0.24' --index-url https://download.pytorch.org/whl/cu121 && \
    pip install --no-cache-dir --force https://modelscope.oss-cn-beijing.aliyuncs.com/packages/tinycudann-1.7-cp310-cp310-linux_x86_64.whl && \
    pip uninstall -y torch-scatter && TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.0;8.6;8.9;9.0" pip install --no-cache-dir -U torch-scatter && \
    pip install --no-cache-dir -U flash_attn vllm; \
    else \
        echo 'cpu unsupport vllm auto-gptq'; \
    fi

COPY requirements /var/modelscope
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/audio.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/cv.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/multi-modal.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/nlp.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/science.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/tests.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
    pip install --no-cache-dir -r /var/modelscope/svr.txt && \
    pip cache purge

COPY examples /modelscope/examples
ENV SETUPTOOLS_USE_DISTUTILS=stdlib
ENV VLLM_USE_MODELSCOPE=True
