fix: repair Dockerfile for ROCm

With these changes, the Docker image can be built and executed
successfully on hosts with AMD devices with ROCm acceleration.
Previously, a ROCm-enabled version of torch would be installed, but
later removed during installation of InvokeAI itself. This was caused by
InvokeAI needing a newer torch version than was previously installed.

The fix consists of multiple components:
* Update the hardcoded versions of torch and torchvision to the versions
  currently used in pyproject.toml, so that a new version need not be
  installed during installation of InvokeAI.
* Specify --extra-index-url on installation of InvokeAI so that even if
  a verison mismatch occurs, the correct torch version should still be
  installed. This also necessitates changing --index-url to
  --extra-index-url for the Torch repo. Otherwise non-torch dependencies
  would not be found.
* In run.sh, build the image for the selected service.
This commit is contained in:
Copper Phosphate 2024-02-13 22:01:30 +01:00 committed by Kent Keirsey
parent 5e77f0d93b
commit c5aeb36230
2 changed files with 5 additions and 5 deletions

View File

@ -18,8 +18,8 @@ ENV INVOKEAI_SRC=/opt/invokeai
ENV VIRTUAL_ENV=/opt/venv/invokeai
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
ARG TORCH_VERSION=2.1.0
ARG TORCHVISION_VERSION=0.16
ARG TORCH_VERSION=2.1.2
ARG TORCHVISION_VERSION=0.16.2
ARG GPU_DRIVER=cuda
ARG TARGETPLATFORM="linux/amd64"
# unused but available
@ -35,7 +35,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then \
extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cpu"; \
elif [ "$GPU_DRIVER" = "rocm" ]; then \
extra_index_url_arg="--index-url https://download.pytorch.org/whl/rocm5.6"; \
extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/rocm5.6"; \
else \
extra_index_url_arg="--extra-index-url https://download.pytorch.org/whl/cu121"; \
fi &&\
@ -54,7 +54,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$GPU_DRIVER" = "cuda" ] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
pip install -e ".[xformers]"; \
else \
pip install -e "."; \
pip install $extra_index_url_arg -e "."; \
fi
# #### Build the Web UI ------------------------------------

View File

@ -21,7 +21,7 @@ run() {
printf "%s\n" "$build_args"
fi
docker compose build $build_args
docker compose build $build_args $service_name
unset build_args
printf "%s\n" "starting service $service_name"