ARG BASE
FROM ${BASE}

ARG GRAPHIC_VERSION
RUN env FORCE_CUDA=1 CUDA_VISIBLE_DEVICES=0 TORCH_CUDA_ARCH_LIST=${GRAPHIC_VERSION} \
  pip install --no-cache-dir git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

CMD ["--listen", "--port", "8080", "--xformers"]
