From b031d1f9d55f2cf76a8d5fc48e532dd51bd7deb0 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Wed, 6 Dec 2023 07:32:36 +0000 Subject: [PATCH] Add method to easily bake models and nodes into the image --- README.md | 22 +++- .../ai-dock/storage_monitor/etc/mappings.sh | 11 +- .../opt/ai-dock/bin/build/layer1/init.sh | 117 ++++++++++++++++++ .../models/diffusers/.gitkeep | 0 .../models/embeddings/.gitkeep | 0 .../stable_diffusion/models/gligen/.gitkeep | 0 .../models/hypernetworks/.gitkeep | 0 .../models/style_models/.gitkeep | 0 .../stable_diffusion/models/unet/.gitkeep | 0 .../models/vae_approx/.gitkeep | 0 build/Dockerfile | 3 +- 11 files changed, 144 insertions(+), 9 deletions(-) create mode 100755 build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/diffusers/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/embeddings/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/gligen/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/hypernetworks/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/style_models/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/unet/.gitkeep create mode 100644 build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/vae_approx/.gitkeep diff --git a/README.md b/README.md index 4ad08b5d..4ba98bbc 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ Tags follow these patterns: Browse [here](https://github.com/ai-dock/comfyui/pkgs/container/comfyui) for an image suitable for your target environment. -You can also self-build from source by editing `.env` and running `docker compose build`. +You can also [build from source](#building-images) by editing `.env` and running `docker compose build`. Supported Python versions: `3.10` @@ -93,15 +93,25 @@ Supported Platforms: `NVIDIA CUDA`, `AMD ROCm`, `CPU` ## Building Images -You can self-build from source by editing `docker-compose.yaml` or `.env` and running `docker compose build`. +You can self-build from source by editing `docker-compose.yaml` or `.env` and running `docker compose build`. -It is a good idea to leave the source tree alone and copy any edits you would like to make into `build/COPY_ROOT_EXTRA/...`. The structure within this directory will be overlayed on `/` at the end of the build process. +It is a good idea to leave the main source tree alone and copy any extra files you would like in the container into `build/COPY_ROOT_EXTRA/...`. The structure within this directory will be overlayed on `/` near the end of the build process. -As this overlaying happens after the main build, it is easy to add extra files such as ML models and datasets to your images. You will also be able to rebuild quickly if your file overrides are made here. +After copying has been completed, the script `build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh` will be executed. A template for this file capable of downloading models and nodes is provided for convenience. -Any directories and files that you add into `opt/storage` will be made available in the running container at `$WORKSPACE/storage`. +Any directories and files that you add into `opt/storage` will be made available in the running container at `$WORKSPACE/storage` through symbolic links. -This directory is monitored by `inotifywait`. Any items appearing in this directory will be automatically linked to the application directories as defined in `/opt/ai-dock/storage_monitor/etc/mappings.sh`. This is particularly useful if you need to run several applications that each need to make use of the stored files. +This directory is monitored by `inotifywait`. Any items appearing here will be automatically symlinked to the application directories as defined in `/opt/ai-dock/storage_monitor/etc/mappings.sh`. + +### Recommended workflow + +- Fork this repository and clone +- Create and switch to a new branch +- Create `.env` to override the `IMAGE_TAG` +- Copy non-public models to `build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/ckpt/` +- Edit `build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh` to download public models and nodes +- Run `docker compose build` +- Run `docker compose push` ## Run Locally diff --git a/build/COPY_ROOT/opt/ai-dock/storage_monitor/etc/mappings.sh b/build/COPY_ROOT/opt/ai-dock/storage_monitor/etc/mappings.sh index 9331cde9..d7c47594 100644 --- a/build/COPY_ROOT/opt/ai-dock/storage_monitor/etc/mappings.sh +++ b/build/COPY_ROOT/opt/ai-dock/storage_monitor/etc/mappings.sh @@ -2,9 +2,16 @@ declare -A storage_map storage_map["stable_diffusion/models/ckpt"]="/opt/ComfyUI/models/checkpoints" -storage_map["stable_diffusion/models/lora"]="/opt/ComfyUI/models/loras" storage_map["stable_diffusion/models/controlnet"]="/opt/ComfyUI/models/controlnet" -storage_map["stable_diffusion/models/vae"]="/opt/ComfyUI/models/vae" +storage_map["stable_diffusion/models/diffusers"]="/opt/ComfyUI/models/diffusers" +storage_map["stable_diffusion/models/embeddings"]="/opt/ComfyUI/models/embeddings" storage_map["stable_diffusion/models/esrgan"]="/opt/ComfyUI/models/upscale_models" +storage_map["stable_diffusion/models/gligen"]="/opt/ComfyUI/models/gligen" +storage_map["stable_diffusion/models/hypernetworks"]="/opt/ComfyUI/models/hypernetworks" +storage_map["stable_diffusion/models/lora"]="/opt/ComfyUI/models/loras" +storage_map["stable_diffusion/models/style_models"]="/opt/ComfyUI/models/style_models" +storage_map["stable_diffusion/models/unet"]="/opt/ComfyUI/models/unet" +storage_map["stable_diffusion/models/vae"]="/opt/ComfyUI/models/vae" +storage_map["stable_diffusion/models/vae_approx"]="/opt/ComfyUI/models/upscale_models" # Add more mappings for other repository directories as needed \ No newline at end of file diff --git a/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh b/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh new file mode 100755 index 00000000..7bd5e536 --- /dev/null +++ b/build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh @@ -0,0 +1,117 @@ +#!/bin/bash + +# Use this layer to add nodes and models + +NODES=( + #"https://github.com/ltdrdata/ComfyUI-Manager" +) + +CHECKPOINT_MODELS=( + #"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt" + #"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.ckpt" + #"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors" + #"https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors" +) + +LORA_MODELS=( + #"https://civitai.com/api/download/models/16576" +) + +VAE_MODELS=( + #"https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.safetensors" + #"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors" + #"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors" +) + +ESRGAN_MODELS=( + #"https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth" + #"https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth" + #"https://huggingface.co/Akumetsu971/SD_Anime_Futuristic_Armor/resolve/main/4x_NMKD-Siax_200k.pth" +) + +CONTROLNET_MODELS=( + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_canny-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_depth-fp16.safetensors" + #"https://huggingface.co/kohya-ss/ControlNet-diff-modules/resolve/main/diff_control_sd15_depth_fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_hed-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_mlsd-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_normal-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_openpose-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_scribble-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_seg-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_canny-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_color-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_depth-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_keypose-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_openpose-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_seg-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_sketch-fp16.safetensors" + #"https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/t2iadapter_style-fp16.safetensors" +) + +### DO NOT EDIT BELOW HERE UNLESS YOU KNOW WHAT YOU ARE DOING ### + +function build_extra_start() { + build_extra_get_nodes + build_extra_get_models \ + "/opt/storage/stable_diffusion/models/ckpt" \ + "${CHECKPOINT_MODELS[@]}" + build_extra_get_models \ + "/opt/storage/stable_diffusion/models/lora" \ + "${LORA_MODELS[@]}" + build_extra_get_models \ + "/opt/storage/stable_diffusion/models/controlnet" \ + "${CONTROLNET_MODELS[@]}" + build_extra_get_models \ + "/opt/storage/stable_diffusion/models/vae" \ + "${VAE_MODELS[@]}" + build_extra_get_models \ + "/opt/storage/stable_diffusion/models/esrgan" \ + "${ESRGAN_MODELS[@]}" +} + +function build_extra_get_nodes() { + for repo in "${NODES[@]}"; do + dir="${repo##*/}" + path="/opt/ComfyUI/custom_nodes/${dir}" + requirements="${path}/requirements.txt" + if [[ -d $path ]]; then + if [[ ${AUTO_UPDATE,,} != "false" ]]; then + printf "Updating node: %s...\n" "${repo}" + ( cd "$path" && git pull ) + if [[ -e $requirements ]]; then + micromamba -n comfyui run ${PIP_INSTALL} -r "$requirements" + fi + fi + else + printf "Downloading node: %s...\n" "${repo}" + git clone "${repo}" "${path}" + if [[ -e $requirements ]]; then + micromamba -n comfyui run ${PIP_INSTALL} -r "${requirements}" + fi + fi + done +} + +function build_extra_get_models() { + if [[ -n $2 ]]; then + dir="$1" + mkdir -p "$dir" + shift + arr=("$@") + + printf "Downloading %s model(s) to %s...\n" "${#arr[@]}" "$dir" + for url in "${arr[@]}"; do + printf "Downloading: %s\n" "${url}" + build_extra_download "${url}" "${dir}" + printf "\n" + done + fi +} + +# Download from $1 URL to $2 file path +function build_extra_download() { + wget -qnc --content-disposition --show-progress -e dotbytes="${3:-4M}" -P "$2" "$1" +} + +build_extra_start \ No newline at end of file diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/diffusers/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/diffusers/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/embeddings/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/embeddings/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/gligen/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/gligen/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/hypernetworks/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/hypernetworks/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/style_models/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/style_models/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/unet/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/unet/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/vae_approx/.gitkeep b/build/COPY_ROOT_EXTRA/opt/storage/stable_diffusion/models/vae_approx/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/build/Dockerfile b/build/Dockerfile index 0e0169fc..57e529b6 100644 --- a/build/Dockerfile +++ b/build/Dockerfile @@ -24,8 +24,9 @@ ENV OPT_SYNC=ComfyUI:serverless:$OPT_SYNC ENV MAMBA_DEFAULT_ENV=comfyui ENV MAMBA_DEFAULT_RUN="micromamba run -n $MAMBA_DEFAULT_ENV" -# Copy overrides and models into a final layer for fast rebuilds +# Copy overrides and models into later layers for fast rebuilds COPY ./COPY_ROOT_EXTRA/ / +RUN /opt/ai-dock/bin/build/layer1/init.sh # Keep init.sh as-is and place additional logic in /opt/ai-dock/bin/preflight.sh CMD ["init.sh"]