feat(docker): remove separate pre-installation of PyTorch in the image

This commit is contained in:
Eugene Brodsky 2024-03-20 14:19:24 -04:00
parent c9350f71be
commit 3278497674

View File

@ -18,8 +18,6 @@ ENV INVOKEAI_SRC=/opt/invokeai
ENV VIRTUAL_ENV=/opt/venv/invokeai ENV VIRTUAL_ENV=/opt/venv/invokeai
ENV PATH="$VIRTUAL_ENV/bin:$PATH" ENV PATH="$VIRTUAL_ENV/bin:$PATH"
ARG TORCH_VERSION=2.1.2
ARG TORCHVISION_VERSION=0.16.2
ARG GPU_DRIVER=cuda ARG GPU_DRIVER=cuda
ARG TARGETPLATFORM="linux/amd64" ARG TARGETPLATFORM="linux/amd64"
# unused but available # unused but available
@ -27,7 +25,12 @@ ARG BUILDPLATFORM
WORKDIR ${INVOKEAI_SRC} WORKDIR ${INVOKEAI_SRC}
# Install pytorch before all other pip packages COPY invokeai ./invokeai
COPY pyproject.toml ./
# Editable mode helps use the same image for development:
# the local working copy can be bind-mounted into the image
# at path defined by ${INVOKEAI_SRC}
# NOTE: there are no pytorch builds for arm64 + cuda, only cpu # NOTE: there are no pytorch builds for arm64 + cuda, only cpu
# x86_64/CUDA is default # x86_64/CUDA is default
RUN --mount=type=cache,target=/root/.cache/pip \ RUN --mount=type=cache,target=/root/.cache/pip \
@ -39,20 +42,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \
else \ else \
extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cu121"; \ extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cu121"; \
fi &&\ fi &&\
pip install $extra_index_url_arg \
torch==$TORCH_VERSION \
torchvision==$TORCHVISION_VERSION
# Install the local package.
# Editable mode helps use the same image for development:
# the local working copy can be bind-mounted into the image
# at path defined by ${INVOKEAI_SRC}
COPY invokeai ./invokeai
COPY pyproject.toml ./
RUN --mount=type=cache,target=/root/.cache/pip \
# xformers + triton fails to install on arm64 # xformers + triton fails to install on arm64
if [ "$GPU_DRIVER" = "cuda" ] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then \ if [ "$GPU_DRIVER" = "cuda" ] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
pip install -e ".[xformers]"; \ pip install $extra_index_url_arg -e ".[xformers]"; \
else \ else \
pip install $extra_index_url_arg -e "."; \ pip install $extra_index_url_arg -e "."; \
fi fi