From 0f3187505bfe76ba24372f135e4f6481097210af Mon Sep 17 00:00:00 2001
From: Self Denial <selfdenial@pm.me>
Date: Tue, 29 Oct 2024 01:36:21 -0600
Subject: [PATCH] feat: Add support for Krita AI Diffusion support

This commit adds support for integrating Krita AI Diffusion models and nodes into the Comfy image editing application. It includes the following changes:

- Downloads recommended models for Krita AI Diffusion if enabled
- Symlinks the downloaded upscale models into the models directory
- Adds environment variables to enable/disable dependancy node packs
- Clones and sets up the required Git repositories for those node packs if enabled
- Installs additional dependencies like aiohttp and tqdm if Krita AI Diffusion is used

These changes allow Comfy to leverage the powerful image generation and editing capabilities provided by the Krita AI Diffusion project. Users can now access advanced features like outpainting, inpainting, upscaling, etc. within the Comfy UI.

The commit also improves the build process by using a cache for pip installs and specifying types for mounted volumes in Docker for better performance and reproducibility.
---
 services/comfy/Dockerfile             | 10 +++++++---
 services/comfy/entrypoint.sh          | 18 +++++++++++++++++-
 services/comfy/extra_model_paths.yaml |  2 ++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/services/comfy/Dockerfile b/services/comfy/Dockerfile
index be55849..b474d34 100644
--- a/services/comfy/Dockerfile
+++ b/services/comfy/Dockerfile
@@ -47,8 +47,12 @@ RUN --mount=type=cache,uid=${USE_UID},gid=${USE_GID},target=${CACHE} \
   git reset --hard ${VERSION}' && \
   pip --cache-dir=${CACHE}/pip install -r requirements.txt && \
   if [ "${USE_KRITA}" = "true" ]; then \
-    USE_CNAUX="true"; USE_IPAPLUS="true"; \
-    USE_INPAINT="true"; USE_TOOLING="true"; \
+    pip --cache-dir=${CACHE}/pip install aiohttp tqdm && \
+    git clone https://github.com/Acly/krita-ai-diffusion.git && \
+    cd krita-ai-diffusion && git checkout main && \
+    git submodule update --init && cd ..; \
+    export USE_CNAUX="true" USE_IPAPLUS="true" \
+    USE_INPAINT="true" USE_TOOLING="true"; \
   fi; \
   if [ "${USE_GGUF}" = "true" ]; then \
     git clone https://github.com/city96/ComfyUI-GGUF.git && \
@@ -74,7 +78,7 @@ RUN --mount=type=cache,uid=${USE_UID},gid=${USE_GID},target=${CACHE} \
   if [ "${USE_INPAINT}" = "true" ]; then \
     git clone https://github.com/Acly/comfyui-inpaint-nodes.git && \
     cd comfyui-inpaint-nodes && git checkout main && \
-    pip install opencv-python && cd ..; \
+    pip --cache-dir=${CACHE}/pip install opencv-python && cd ..; \
   fi; \
   if [ "${USE_TOOLING}" = "true" ]; then \
     git clone https://github.com/Acly/comfyui-tooling-nodes.git && \
diff --git a/services/comfy/entrypoint.sh b/services/comfy/entrypoint.sh
index f4ffe68..53cb6ce 100755
--- a/services/comfy/entrypoint.sh
+++ b/services/comfy/entrypoint.sh
@@ -23,6 +23,14 @@ for to_path in "${!MOUNTS[@]}"; do
   echo Mounted $(basename "${from_path}")
 done
 
+
+if [ "${USE_KRITA}" = "true" ]; then
+  if [ "${KRITA_DOWNLOAD_MODELS:-false}" = "true" ]; then
+    cd "${ROOT}/krita-ai-diffusion/scripts" && python scripts/download_models.py --recommended /data && cd ..
+    cd "${ROOT}/models/" mv -v upscale_models upscale_models.stock && ln -sT /data/models/upscale_models upscale_models
+  fi
+  export USE_CNAUX="true" USE_IPAPLUS="true" USE_INPAINT="true"; USE_TOOLING="true"
+fi
 if [ "${USE_GGUF}" = "true" ]; then
   [ ! -e "${CUSTOM_NODES}/ComfyUI-GGUF" ] && mv "${ROOT}/ComfyUI-GGUF" "${CUSTOM_NODES}"/
 fi
@@ -42,7 +50,15 @@ fi
 if [ "${USE_CNAUX}" = "true" ]; then
   [ ! -e "${CUSTOM_NODES}/comfyui_controlnet_aux" ] && mv "${ROOT}/comfyui_controlnet_aux" "${CUSTOM_NODES}"/
 fi
-
+if [ "${USE_IPAPLUS}" = "true" ]; then
+  [ ! -e "${CUSTOM_NODES}/ComfyUI_IPAdapter_plus" ] && mv "${ROOT}/ComfyUI_IPAdapter_plus" "${CUSTOM_NODES}"/
+fi
+if [ "${USE_INPAINT}" = "true" ]; then
+  [ ! -e "${CUSTOM_NODES}/comfyui-inpaint-nodes" ] && mv "${ROOT}/comfyui-inpaint-nodes" "${CUSTOM_NODES}"/
+fi
+if [ "${USE_TOOLING}" = "true" ]; then
+  [ ! -e "${CUSTOM_NODES}/comfyui-tooling-nodes" ] && mv "${ROOT}/comfyui-tooling-nodes" "${CUSTOM_NODES}"/
+fi
 
 if [ -f "/data/config/comfy/startup.sh" ]; then
   pushd ${ROOT}
diff --git a/services/comfy/extra_model_paths.yaml b/services/comfy/extra_model_paths.yaml
index 2443bdd..365f5d2 100644
--- a/services/comfy/extra_model_paths.yaml
+++ b/services/comfy/extra_model_paths.yaml
@@ -18,6 +18,8 @@ a111:
   unet: models/unet
   clip_vision: models/clip_vision
   xlabs: models/xlabs
+  inpaint: models/inpaint
+  ipadapter: models/ipadapter
 
   custom_nodes: config/comfy/custom_nodes