From eb702a5049f2a15423bab7f5406463234fbf9d14 Mon Sep 17 00:00:00 2001
From: mauwii <Mauwii@outlook.de>
Date: Sun, 5 Feb 2023 02:39:29 +0100
Subject: [PATCH] fix env.sh, update Dockerfile, update build.sh

env.sh:
- move check for torch to CONVTAINER_FLAVOR detection

Dockerfile
- only mount `/var/cache/apt` for apt related steps
- remove `docker-clean` from `/etc/apt/apt.conf.d` for BuildKit cache
- remove apt-get clean for BuildKit cache
- only copy frontend to frontend-builder
- mount `/usr/local/share/.cache/yarn` in frountend-builder
- separate steps for yarn install and yarn build
- build pytorch in pyproject-builder

build.sh
- prepare for installation with extras
---
 docker/Dockerfile | 56 +++++++++++++++++++----------------------------
 docker/build.sh   |  1 +
 docker/env.sh     | 40 ++++++++++++++++-----------------
 3 files changed, 44 insertions(+), 53 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 475f8dc55e..a3a33112e5 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -4,21 +4,22 @@
 
 ARG PYTHON_VERSION=3.9
 ##################
-### base image ###
+##  base image  ##
 ##################
 FROM python:${PYTHON_VERSION}-slim AS python-base
 
+# prepare for buildkit cache
+RUN rm -f /etc/apt/apt.conf.d/docker-clean
+
 # Install necesarry packages
 RUN \
   --mount=type=cache,target=/var/cache/apt,sharing=locked \
-  --mount=type=cache,target=/var/lib/apt,sharing=locked \
   apt-get update \
   && apt-get install -y \
     --no-install-recommends \
     libgl1-mesa-glx=20.3.* \
     libglib2.0-0=2.66.* \
     libopencv-dev=4.5.* \
-  && apt-get clean \
   && rm -rf /var/lib/apt/lists/*
 
 # set working directory and path
@@ -28,41 +29,43 @@ WORKDIR ${APPDIR}
 ENV PATH=${APPDIR}/${APPNAME}/bin:$PATH
 
 ######################
-### build frontend ###
+##  build frontend  ##
 ######################
 FROM node:lts as frontend-builder
 
 # Copy Sources
 ARG APPDIR=/usr/src
 WORKDIR ${APPDIR}
-COPY --link . .
+COPY ./invokeai/frontend ./invokeai/frontend
 
-# install dependencies and build frontend
+# install dependencies
 WORKDIR ${APPDIR}/invokeai/frontend
 RUN \
-  --mount=type=cache,target=/usr/local/share/.cache/yarn/v6 \
+  --mount=type=cache,target=/usr/local/share/.cache/yarn \
   yarn install \
   --prefer-offline \
   --frozen-lockfile \
   --non-interactive \
-  --production=false \
-  && yarn build
+  --production=false
+
+# build frontend
+RUN yarn build
 
 ###################################
-### install python dependencies ###
+##  install python dependencies  ##
 ###################################
 FROM python-base AS pyproject-builder
+ENV PIP_USE_PEP517=1
 
 # Install dependencies
 RUN \
   --mount=type=cache,target=/var/cache/apt,sharing=locked \
-  --mount=type=cache,target=/var/lib/apt,sharing=locked \
   apt-get update \
   && apt-get install -y \
     --no-install-recommends \
+    build-essential=12.9 \
     gcc=4:10.2.* \
     python3-dev=3.9.* \
-  && apt-get clean \
   && rm -rf /var/lib/apt/lists/*
 
 # create virtual environment
@@ -70,18 +73,21 @@ RUN python3 -m venv "${APPNAME}" \
   --upgrade-deps
 
 # copy sources
-COPY --from=frontend-builder ${APPDIR} .
+COPY --link . .
+COPY --from=frontend-builder ${APPDIR}/invokeai/frontend/dist ${APPDIR}/invokeai/frontend/dist
 
 # install pyproject.toml
 ARG PIP_EXTRA_INDEX_URL
 ENV PIP_EXTRA_INDEX_URL ${PIP_EXTRA_INDEX_URL}
+ARG PIP_PACKAGE=.
 RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked \
-  "${APPDIR}/${APPNAME}/bin/pip" install \
-  --use-pep517 \
-  .
+  "${APPDIR}/${APPNAME}/bin/pip" install ${PIP_PACKAGE}
+
+# build patchmatch
+RUN python3 -c "from patchmatch import patch_match"
 
 #####################
-### runtime image ###
+##  runtime image  ##
 #####################
 FROM python-base AS runtime
 
@@ -90,22 +96,6 @@ COPY --from=pyproject-builder ${APPDIR}/${APPNAME} ${APPDIR}/${APPNAME}
 ENV INVOKEAI_ROOT=/data
 ENV INVOKE_MODEL_RECONFIGURE="--yes --default_only"
 
-# build patchmatch
-RUN \
-  --mount=type=cache,target=/var/cache/apt,sharing=locked \
-  --mount=type=cache,target=/var/lib/apt,sharing=locked \
-  apt-get update \
-  && apt-get install -y \
-    --no-install-recommends \
-    build-essential=12.9 \
-  && PYTHONDONTWRITEBYTECODE=1 \
-    python3 -c "from patchmatch import patch_match" \
-  && apt-get remove -y \
-    --autoremove \
-    build-essential \
-  && apt-get autoclean \
-  && rm -rf /var/lib/apt/lists/*
-
 # set Entrypoint and default CMD
 ENTRYPOINT [ "invokeai" ]
 CMD [ "--web", "--host=0.0.0.0" ]
diff --git a/docker/build.sh b/docker/build.sh
index c5670aa29b..bed4d8a113 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -39,5 +39,6 @@ DOCKER_BUILDKIT=1 docker build \
     --platform="${PLATFORM}" \
     --tag="${CONTAINER_IMAGE}" \
     ${PIP_EXTRA_INDEX_URL:+--build-arg="PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}"} \
+    ${PIP_PACKAGE:+--build-arg="PIP_PACKAGE=${PIP_PACKAGE}"} \
     --file="${DOCKERFILE}" \
     ..
diff --git a/docker/env.sh b/docker/env.sh
index 6987e12fb8..d6b0699ce5 100644
--- a/docker/env.sh
+++ b/docker/env.sh
@@ -1,27 +1,27 @@
 #!/usr/bin/env bash
 
-if python -c "import torch" &>/dev/null; then
-  if [[ -z "$PIP_EXTRA_INDEX_URL" ]]; then
-    # Decide which container flavor to build if not specified
-    if [[ -z "$CONTAINER_FLAVOR" ]]; then
-      # Check for CUDA and ROCm
-      CUDA_AVAILABLE=$(python -c "import torch;print(torch.cuda.is_available())")
-      ROCM_AVAILABLE=$(python -c "import torch;print(torch.version.hip is not None)")
-      if [[ "$(uname -s)" != "Darwin" && "${CUDA_AVAILABLE}" == "True" ]]; then
-        CONTAINER_FLAVOR="cuda"
-      elif [[ "$(uname -s)" != "Darwin" && "${ROCM_AVAILABLE}" == "True" ]]; then
-        CONTAINER_FLAVOR="rocm"
-      else
-        CONTAINER_FLAVOR="cpu"
-      fi
-    fi
-    # Set PIP_EXTRA_INDEX_URL based on container flavor
-    if [[ "$CONTAINER_FLAVOR" == "rocm" ]]; then
-      PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/rocm"
-    elif [[ "$CONTAINER_FLAVOR" == "cpu" ]]; then
-      PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
+if [[ -z "$PIP_EXTRA_INDEX_URL" ]]; then
+  # Decide which container flavor to build if not specified
+  if [[ -z "$CONTAINER_FLAVOR" ]] && python -c "import torch" &>/dev/null; then
+    # Check for CUDA and ROCm
+    CUDA_AVAILABLE=$(python -c "import torch;print(torch.cuda.is_available())")
+    ROCM_AVAILABLE=$(python -c "import torch;print(torch.version.hip is not None)")
+    if [[ "$(uname -s)" != "Darwin" && "${CUDA_AVAILABLE}" == "True" ]]; then
+      CONTAINER_FLAVOR="cuda"
+    elif [[ "$(uname -s)" != "Darwin" && "${ROCM_AVAILABLE}" == "True" ]]; then
+      CONTAINER_FLAVOR="rocm"
+    else
+      CONTAINER_FLAVOR="cpu"
     fi
   fi
+  # Set PIP_EXTRA_INDEX_URL based on container flavor
+  if [[ "$CONTAINER_FLAVOR" == "rocm" ]]; then
+    PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/rocm"
+  elif [[ "$CONTAINER_FLAVOR" == "cpu" ]]; then
+    PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
+  # elif [[ -z "$CONTAINER_FLAVOR" || "$CONTAINER_FLAVOR" == "cuda" ]]; then
+  #   PIP_PACKAGE=${PIP_PACKAGE-".[xformers]"}
+  fi
 fi
 
 # Variables shared by build.sh and run.sh