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

@@ -1,29 +1,30 @@
+FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
+
 # Limited system user UID
 ARG USE_UID=991
 # Limited system user GID
 ARG USE_GID=991
 # Latest tag or bleeding edge commit
 ARG USE_EDGE=false
-# ComfyUI-GGUF support
+# ComfyUI-GGUF
 ARG USE_GGUF=false
-# x-flux-comfyui support
+# x-flux-comfyui
 ARG USE_XFLUX=false
-# comfyui_controlnet_aux support
+# comfyui_controlnet_aux
 ARG USE_CNAUX=false
-
-FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
-
-# Use args
-ARG USE_UID
-ARG USE_GID
-ARG USE_EDGE
-ARG USE_GGUF
-ARG USE_XFLUX
-ARG USE_CNAUX
+# krita-ai-diffusion
+ARG USE_KRITA=false
+# ComfyUI_IPAdapter_plus
+ARG USE_IPAPLUS=false
+# comfyui-inpaint-nodes
+ARG USE_INPAINT=false
+# comfyui-tooling-nodes
+ARG USE_TOOLING=false
 
 ENV DEBIAN_FRONTEND=noninteractive PIP_PREFER_BINARY=1 USE_EDGE=$USE_EDGE
 ENV USE_GGUF=$USE_GGUF USE_XFLUX=$USE_XFLUX ROOT=/stable-diffusion
-ENV CACHE=/home/app/.cache USE_CNAUX=$USE_CNAUX
+ENV CACHE=/home/app/.cache USE_CNAUX=$USE_CNAUX USE_KRITA=$USE_KRITA
+ENV USE_IPAPLUS=$USE_IPAPLUS USE_INPAINT=$USE_INPAINT USE_TOOLING=$USE_TOOLING
 
 # User/Group
 RUN groupadd -r app -g ${USE_GID} && useradd --no-log-init -m -r -g app app -u ${USE_UID} && \
@@ -45,6 +46,14 @@ RUN --mount=type=cache,uid=${USE_UID},gid=${USE_GID},target=${CACHE} \
   if [ "${USE_EDGE}" = "true" ]; then VERSION=$(git describe --abbrev=7); fi && \
   git reset --hard ${VERSION}' && \
   pip --cache-dir=${CACHE}/pip install -r requirements.txt && \
+  if [ "${USE_KRITA}" = "true" ]; then \
+    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 && \
     cd ComfyUI-GGUF && git checkout main && \
@@ -59,8 +68,21 @@ RUN --mount=type=cache,uid=${USE_UID},gid=${USE_GID},target=${CACHE} \
     git clone https://github.com/Fannovel16/comfyui_controlnet_aux.git && \
     cd comfyui_controlnet_aux && git checkout main && \
     pip --cache-dir=${CACHE}/pip install -r requirements.txt && \
-    # This extra step to separate onnxruntime installation is required to restore onnx cuda support
+    # This extra step to separate onnxruntime installation is required to restore onnx cuda support \
     pip --cache-dir=${CACHE}/pip install onnxruntime && pip --cache-dir=${CACHE}/pip install onnxruntime-gpu && cd ..; \
+  fi; \
+  if [ "${USE_IPAPLUS}" = "true" ]; then \
+    git clone https://github.com/cubiq/ComfyUI_IPAdapter_plus.git && \
+    cd ComfyUI_IPAdapter_plus && git checkout main && cd ..; \
+  fi; \
+  if [ "${USE_INPAINT}" = "true" ]; then \
+    git clone https://github.com/Acly/comfyui-inpaint-nodes.git && \
+    cd comfyui-inpaint-nodes && git checkout main && \
+    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 && \
+    cd comfyui-tooling-nodes && git checkout main && cd ..; \
   fi
 
 WORKDIR ${ROOT}

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}

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