Skip to content

Commit

Permalink
Add method to easily bake models and nodes into the image
Browse files Browse the repository at this point in the history
  • Loading branch information
robballantyne committed Dec 6, 2023
1 parent a69f6eb commit b031d1f
Show file tree
Hide file tree
Showing 11 changed files with 144 additions and 9 deletions.
22 changes: 16 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`

Expand All @@ -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

Expand Down
11 changes: 9 additions & 2 deletions build/COPY_ROOT/opt/ai-dock/storage_monitor/etc/mappings.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
117 changes: 117 additions & 0 deletions build/COPY_ROOT_EXTRA/opt/ai-dock/bin/build/layer1/init.sh
Original file line number Diff line number Diff line change
@@ -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
Empty file.
Empty file.
Empty file.
Empty file.
Empty file.
Empty file.
Empty file.
3 changes: 2 additions & 1 deletion build/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

0 comments on commit b031d1f

Please sign in to comment.