| ARG CUDA_VERSION=cu12 | |
| FROM openmmlab/lmdeploy:latest-cu12 AS cu12 | |
| ENV CUDA_VERSION_SHORT=cu123 | |
| FROM openmmlab/lmdeploy:latest-cu11 AS cu11 | |
| ENV CUDA_VERSION_SHORT=cu118 | |
| FROM ${CUDA_VERSION} AS final | |
| RUN python3 -m pip install timm!=1.0.23 | |
| RUN python3 -m pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+${CUDA_VERSION_SHORT}torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl | |