From 47508b8d6c78fbe1fdfabe8c6f004a49c015badb Mon Sep 17 00:00:00 2001
From: Heathen711 <heathen711@me.com>
Date: Thu, 3 Jul 2025 06:01:27 +0000
Subject: [PATCH] bugfix(docker) combined the dockerfiles and reduced image
 size

---
 docker/Dockerfile         |  53 ++++++++++++++------
 docker/Dockerfile-rocm    | 102 --------------------------------------
 docker/docker-compose.yml |   4 --
 3 files changed, 39 insertions(+), 120 deletions(-)
 delete mode 100644 docker/Dockerfile-rocm

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 6a34f30fa6..6319dc29b4 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -31,7 +31,8 @@ RUN --mount=type=cache,target=/var/cache/apt \
     libglx-mesa0 \
     build-essential \
     libopencv-dev \
-    libstdc++-10-dev
+    libstdc++-10-dev \
+    wget
 
 ENV \
     PYTHONUNBUFFERED=1 \
@@ -44,7 +45,6 @@ ENV \
     UV_MANAGED_PYTHON=1 \
     UV_LINK_MODE=copy \
     UV_PROJECT_ENVIRONMENT=/opt/venv \
-    UV_INDEX="https://download.pytorch.org/whl/cu124" \
     INVOKEAI_ROOT=/invokeai \
     INVOKEAI_HOST=0.0.0.0 \
     INVOKEAI_PORT=9090 \
@@ -54,6 +54,10 @@ ENV \
 
 ARG GPU_DRIVER=cuda
 
+ARG CUDA_TORCH="https://download.pytorch.org/whl/cu124"
+ARG CPU_TORCH="https://download.pytorch.org/whl/cpu"
+ARG ROCM_TORCH="https://download.pytorch.org/whl/rocm6.2.4"
+
 # Install `uv` for package management
 COPY --from=ghcr.io/astral-sh/uv:0.6.9 /uv /uvx /bin/
 
@@ -72,23 +76,41 @@ WORKDIR ${INVOKEAI_SRC}
 # x86_64/CUDA is the default
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
-    --mount=type=bind,source=uv.lock,target=uv.lock \
+    # Cannot use uv sync and uv.lock as that is locked to CUDA version packages, which breaks rocm...
+    # --mount=type=bind,source=uv.lock,target=uv.lock \
     # this is just to get the package manager to recognize that the project exists, without making changes to the docker layer
     --mount=type=bind,source=invokeai/version,target=invokeai/version \
-    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then UV_INDEX="https://download.pytorch.org/whl/cpu"; \
-    elif [ "$GPU_DRIVER" = "rocm" ]; then UV_INDEX="https://download.pytorch.org/whl/rocm6.2"; \
+    ulimit -n 30000 && \
+    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then export UV_INDEX="$CPU_TORCH"; \
+    elif [ "$GPU_DRIVER" = "rocm" ]; then export UV_INDEX="$ROCM_TORCH"; \
+    else export UV_INDEX="$CUDA_TORCH"; \
     fi && \
-    uv sync --frozen
+    uv venv --python 3.12 && \
+    # Use the public version to install existing known dependencies but using the UV_INDEX, not the hardcoded URLs within the uv.lock
+    uv pip install invokeai
+
+RUN --mount=type=cache,target=/var/cache/apt \
+    --mount=type=cache,target=/var/lib/apt \
+    if [ "$GPU_DRIVER" = "rocm" ]; then \
+    wget -O /tmp/amdgpu-install.deb \
+    https://repo.radeon.com/amdgpu-install/6.2.4/ubuntu/noble/amdgpu-install_6.2.60204-1_all.deb && \
+    apt install -y /tmp/amdgpu-install.deb && \
+    apt update && \
+    amdgpu-install --usecase=rocm -y && \
+    apt-get autoclean && \
+    apt clean && \
+    rm -rf /tmp/* /var/tmp/* && \
+    usermod -a -G render ubuntu && \
+    usermod -a -G video ubuntu && \
+    echo "\\n/opt/rocm/lib\\n/opt/rocm/lib64" >> /etc/ld.so.conf.d/rocm.conf && \
+    ldconfig && \
+    update-alternatives --auto rocm; \
+    fi
 
 # build patchmatch
 RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
 RUN python -c "from patchmatch import patch_match"
 
-# Link amdgpu.ids for ROCm builds
-# contributed by https://github.com/Rubonnek
-RUN mkdir -p "/opt/amdgpu/share/libdrm" &&\
-    ln -s "/usr/share/libdrm/amdgpu.ids" "/opt/amdgpu/share/libdrm/amdgpu.ids"
-
 RUN mkdir -p ${INVOKEAI_ROOT} && chown -R ${CONTAINER_UID}:${CONTAINER_GID} ${INVOKEAI_ROOT}
 
 COPY docker/docker-entrypoint.sh ./
@@ -105,9 +127,12 @@ COPY invokeai ${INVOKEAI_SRC}/invokeai
 # in a previous layer
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
-    --mount=type=bind,source=uv.lock,target=uv.lock \
-    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then UV_INDEX="https://download.pytorch.org/whl/cpu"; \
-    elif [ "$GPU_DRIVER" = "rocm" ]; then UV_INDEX="https://download.pytorch.org/whl/rocm6.2"; \
+    # Cannot use the uv.lock as that is locked to CUDA version packages, which breaks rocm...
+    # --mount=type=bind,source=uv.lock,target=uv.lock \
+    ulimit -n 30000 && \
+    if [ "$TARGETPLATFORM" = "linux/arm64" ] || [ "$GPU_DRIVER" = "cpu" ]; then export UV_INDEX="$CPU_TORCH"; \
+    elif [ "$GPU_DRIVER" = "rocm" ]; then export UV_INDEX="$ROCM_TORCH"; \
+    else export UV_INDEX="$CUDA_TORCH"; \
     fi && \
     uv pip install -e .
 
diff --git a/docker/Dockerfile-rocm b/docker/Dockerfile-rocm
deleted file mode 100644
index b6599663d0..0000000000
--- a/docker/Dockerfile-rocm
+++ /dev/null
@@ -1,102 +0,0 @@
-# syntax=docker/dockerfile:1.4
-
-#### Web UI ------------------------------------
-
-FROM docker.io/node:22-slim AS web-builder
-ENV PNPM_HOME="/pnpm"
-ENV PATH="$PNPM_HOME:$PATH"
-RUN corepack use pnpm@8.x
-RUN corepack enable
-
-WORKDIR /build
-COPY invokeai/frontend/web/ ./
-RUN pnpm install --frozen-lockfile
-RUN npx vite build
-
-## Backend ---------------------------------------
-
-FROM library/ubuntu:24.04
-
-ARG DEBIAN_FRONTEND=noninteractive
-RUN rm -f /etc/apt/apt.conf.d/docker-clean; echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache
-RUN apt update && apt install -y --no-install-recommends \
-    ca-certificates \
-    git \
-    gosu \
-    libglib2.0-0 \
-    libgl1 \
-    libglx-mesa0 \
-    build-essential \
-    libopencv-dev \
-    libstdc++-10-dev \
-    wget
-
-ENV \
-    PYTHONUNBUFFERED=1 \
-    PYTHONDONTWRITEBYTECODE=1 \
-    VIRTUAL_ENV=/opt/venv \
-    INVOKEAI_SRC=/opt/invokeai \
-    PYTHON_VERSION=3.12 \
-    UV_PYTHON=3.12 \
-    UV_COMPILE_BYTECODE=1 \
-    UV_MANAGED_PYTHON=1 \
-    UV_LINK_MODE=copy \
-    UV_PROJECT_ENVIRONMENT=/opt/venv \
-    UV_INDEX="https://download.pytorch.org/whl/rocm6.2.4" \
-    INVOKEAI_ROOT=/invokeai \
-    INVOKEAI_HOST=0.0.0.0 \
-    INVOKEAI_PORT=9090 \
-    PATH="/opt/venv/bin:$PATH" \
-    CONTAINER_UID=${CONTAINER_UID:-1000} \
-    CONTAINER_GID=${CONTAINER_GID:-1000}
-
-RUN wget -O /tmp/amdgpu-install.deb \
-    https://repo.radeon.com/amdgpu-install/6.2.4/ubuntu/noble/amdgpu-install_6.2.60204-1_all.deb && \
-    apt install -y /tmp/amdgpu-install.deb && \
-    apt update
-
-RUN amdgpu-install rocm -y
-RUN usermod -a -G render ubuntu && usermod -a -G video ubuntu
-RUN echo "\\n/opt/rocm/lib\\n/opt/rocm/lib64" >> /etc/ld.so.conf.d/rocm.conf && ldconfig
-RUN update-alternatives --auto rocm
-
-# Install `uv` for package management
-COPY --from=ghcr.io/astral-sh/uv:0.6.9 /uv /uvx /bin/
-
-# Install python & allow non-root user to use it by traversing the /root dir without read permissions
-RUN uv python install ${PYTHON_VERSION} && \
-    # chmod --recursive a+rX /root/.local/share/uv/python
-    chmod 711 /root
-
-WORKDIR ${INVOKEAI_SRC}
-
-# Install project's dependencies as a separate layer so they aren't rebuilt every commit.
-# bind-mount instead of copy to defer adding sources to the image until next layer.
-RUN --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
-    # this is just to get the package manager to recognize that the project exists, without making changes to the docker layer
-    --mount=type=bind,source=invokeai/version,target=invokeai/version \
-    ulimit -n 30000 && \
-    uv sync
-
-# build patchmatch
-RUN cd /usr/lib/$(uname -p)-linux-gnu/pkgconfig/ && ln -sf opencv4.pc opencv.pc
-RUN python -c "from patchmatch import patch_match"
-
-RUN mkdir -p ${INVOKEAI_ROOT} && chown -R ${CONTAINER_UID}:${CONTAINER_GID} ${INVOKEAI_ROOT}
-
-COPY docker/docker-entrypoint.sh ./
-ENTRYPOINT ["/opt/invokeai/docker-entrypoint.sh"]
-CMD ["invokeai-web"]
-
-# --link requires buldkit w/ dockerfile syntax 1.4, does not work with podman
-COPY --link --from=web-builder /build/dist ${INVOKEAI_SRC}/invokeai/frontend/web/dist
-
-# add sources last to minimize image changes on code changes
-COPY invokeai ${INVOKEAI_SRC}/invokeai
-
-# this should not increase image size because we've already installed dependencies
-# in a previous layer
-RUN --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
-    ulimit -n 30000 && \
-    uv pip install -e .
-
diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
index a02c1727c8..0c2b6fbd78 100644
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -47,10 +47,6 @@ services:
 
   invokeai-rocm:
     <<: *invokeai
-    image: "ghcr.io/invoke-ai/invokeai:latest-rocm"
-    build:
-      context: ..
-      dockerfile: docker/Dockerfile-rocm
     devices:
       - /dev/kfd:/dev/kfd
       - /dev/dri:/dev/dri