From 33f01eb0df91692824889a8affc61b8c6873a5e6 Mon Sep 17 00:00:00 2001
From: Bernhard Frauendienst <git@nospam.obeliks.de>
Date: Thu, 7 Sep 2023 21:43:55 +0200
Subject: [PATCH] Add Dockerfile

Based on ComfyUI PR #530

Co-Authored-By: ZacharyACoon <show>
---
 .dockerignore       |  8 +++++
 Dockerfile          | 78 +++++++++++++++++++++++++++++++++++++++++++++
 README.md           | 23 +++++++++++++
 docker-compose.yaml | 21 ++++++++++++
 4 files changed, 130 insertions(+)
 create mode 100644 .dockerignore
 create mode 100644 Dockerfile
 create mode 100644 docker-compose.yaml

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 000000000..827e02326
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,8 @@
+.*
+!.git
+__pycache__/
+*.py[cod]
+input
+models
+notebooks
+output
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 000000000..d6b891068
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,78 @@
+# syntax=docker/dockerfile:1
+
+ARG PYTHON_VERSION=3.12
+
+FROM python:${PYTHON_VERSION}-slim
+
+ARG PYTORCH_INSTALL_ARGS=""
+ARG EXTRA_ARGS=""
+ARG USERNAME=comfyui
+ARG USER_UID=1000
+ARG USER_GID=${USER_UID}
+
+# Fail fast on errors or unset variables
+SHELL ["/bin/bash", "-eux", "-o", "pipefail", "-c"]
+
+RUN <<EOF
+	groupadd --gid ${USER_GID} ${USERNAME}
+	useradd --uid ${USER_UID} --gid ${USER_GID} -m ${USERNAME}
+EOF
+
+RUN <<EOF
+	apt-get update
+	apt-get install -y --no-install-recommends \
+		git \
+		git-lfs \
+		rsync \
+		fonts-recommended
+EOF
+
+# run instructions as user
+USER ${USER_UID}:${USER_GID}
+
+WORKDIR /app
+
+ENV XDG_CACHE_HOME=/cache
+ENV PIP_CACHE_DIR=/cache/pip
+ENV VIRTUAL_ENV=/app/venv
+ENV VIRTUAL_ENV_CUSTOM=/app/custom_venv
+
+# create cache directory. During build we will use a cache mount,
+# but later this is useful for custom node installs
+RUN --mount=type=cache,target=/cache/,uid=${USER_UID},gid=${USER_GID} \
+	mkdir -p ${PIP_CACHE_DIR}
+
+# create virtual environment to manage packages
+RUN python -m venv ${VIRTUAL_ENV}
+
+# run python from venv (prefer custom_venv over baked-in one)
+ENV PATH="${VIRTUAL_ENV_CUSTOM}/bin:${VIRTUAL_ENV}/bin:${PATH}"
+
+RUN --mount=type=cache,target=/cache/,uid=${USER_UID},gid=${USER_GID} \
+	pip install torch torchvision torchaudio ${PYTORCH_INSTALL_ARGS}
+
+# copy requirements files first so packages can be cached separately
+COPY --chown=${USER_UID}:${USER_GID} requirements.txt .
+RUN --mount=type=cache,target=/cache/,uid=${USER_UID},gid=${USER_GID} \
+	pip install -r requirements.txt
+
+# Not strictly required for comfyui, but prevents non-working variants of
+# cv2 being pulled in by custom nodes
+RUN --mount=type=cache,target=/cache/,uid=${USER_UID},gid=${USER_GID} \
+	pip install opencv-python-headless
+
+COPY --chown=${USER_UID}:${USER_GID} . .
+
+COPY --chown=nobody:${USER_GID} .git .git
+
+# default environment variables
+ENV COMFYUI_ADDRESS=0.0.0.0
+ENV COMFYUI_PORT=8188
+ENV COMFYUI_EXTRA_ARGS=""
+# default start command
+CMD \
+	if [ -d "${VIRTUAL_ENV_CUSTOM}" ]; then \
+		rsync -aP "${VIRTUAL_ENV}/" "${VIRTUAL_ENV_CUSTOM}/" ;\
+		sed -i "s!${VIRTUAL_ENV}!${VIRTUAL_ENV_CUSTOM}!g" "${VIRTUAL_ENV_CUSTOM}/pyvenv.cfg" ;\
+	fi ;\
+	python -u main.py --listen "${COMFYUI_ADDRESS}" --port "${COMFYUI_PORT}" ${EXTRA_ARGS} ${COMFYUI_EXTRA_ARGS}
diff --git a/README.md b/README.md
index a807ea9d6..bf2da4151 100644
--- a/README.md
+++ b/README.md
@@ -235,6 +235,29 @@ Install the dependencies by opening your terminal inside the ComfyUI folder and:
 
 After this you should have everything installed and can proceed to running ComfyUI.
 
+
+## Docker
+
+You can build a docker image with the `Dockerfile` in this repo.
+Specify, `PYTORCH_INSTALL_ARGS` build arg with one of the PyTorch commands above to build for AMD or NVIDIA GPUs.
+```shell
+docker build --build-arg PYTORCH_INSTALL_ARGS="--index-url https://download.pytorch.org/whl/cu122" .
+```
+```shell
+docker build --build-arg PYTORCH_INSTALL_ARGS="--index-url https://download.pytorch.org/whl/rocm6.2" .
+```
+The `Dockerfile` requires BuildKit to be enabled. If your Docker does not support the buildx command, you can
+enable BuildKit by setting the `DOCKER_BUILDKIT` environment variable.
+```shell
+DOCKER_BUILDKIT=1 docker build --build-arg PYTORCH_INSTALL_ARGS="--index-url https://download.pytorch.org/whl/cu124" .
+```
+> [!NOTE]
+> For building the CPU-only image, it is recommended that you add the --cpu flag to the EXTRA_ARGS build arg:
+>
+> ```shell
+> docker build --build-arg PYTORCH_INSTALL_ARGS="--index-url https://download.pytorch.org/whl/cpu" --build-arg EXTRA_ARGS=--cpu .
+> ```
+
 ### Others:
 
 #### Apple Mac silicon
diff --git a/docker-compose.yaml b/docker-compose.yaml
new file mode 100644
index 000000000..858840fbe
--- /dev/null
+++ b/docker-compose.yaml
@@ -0,0 +1,21 @@
+services:
+  comfyui:
+    user: "1000:1000"
+    build: .
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    ports:
+      - "8188:8188"
+    volumes:
+      - "./models:/app/models"
+      - "./input:/app/input"
+      - "./temp:/app/output/temp"
+      - "./output:/app/output"
+      - "./user:/app/user"
+      - "./custom_venv:/app/custom_venv"
+      - "./custom_nodes:/app/custom_nodes"