Skip to content

Commit

Permalink
feat: network-volume; execution time config; skip default images; acc…
Browse files Browse the repository at this point in the history
…ess ComfyUI via web (#35)

* feat: provide option to run the handler locally as API

* ci: run the workflow on our extended instance

* feat: the local API should run on 0.0.0.0

* feat: make the image smaller

* ci: use semantic-version to create releases automatically

* chore: we don't want to break anyone with a minor release

* docs: added section for local API testing

* ci: use custom runner

* fix: added .releaserc, otherwise semantic-release will complain about a missing "package.json"

* feat: support network volumes, skip default models (#16)

* Support network volumes

* README tweaks

* docs: added comment on what is happening

* feat: don't overwrite the default paths, but add "runpod_worker_comfy" to have additional paths

* docs: updated "bring your own models"

---------

Co-authored-by: Tim Pietrusky <timpietrusky@gmail.com>

* feat: provide access to ComfyUI via web

* fix: use the full path to the output image

* feat: added env vars COMFY_POLLING_INTERVAL_MS and COMFY_POLLING_MAX_RETRIES

* test: added "subfolder"

* docs: update CUDA guide for Ubuntu

* feat(network-volume): added "custom_nodes"

* chore: removed lora from civitai

* docs: removed lora from civitai

* docs: simplified multiple sections; added an image to describe how to find the endpointID

* feat: remove "custom_nodes"

* feat: move models before python installation

* feat: added "runpod-volume"; use "dev" instead of "latest"

* docs: removed "custom nodes"; fixed links to GitHub Actions

* docs: fixed typo in GitHub

* docs: fixed hosts for local testing

* docs: updated link to docs from RunPod to create a network volume

* feat: expose the port of ComfyUI

* docs: fixed example link to download sdxl-turbo; removed nodes from network volume

* docs: moved example response into their own block

* chore: added example workflows for webp and sdxl-turbo

---------

Co-authored-by: Meptl <git-client@meptl.com>
  • Loading branch information
TimPietrusky and Meptl authored Jun 4, 2024
1 parent 15bd612 commit 070cde5
Show file tree
Hide file tree
Showing 10 changed files with 422 additions and 87 deletions.
16 changes: 8 additions & 8 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,12 @@ RUN git clone https://github.com/comfyanonymous/ComfyUI.git /comfyui
# Change working directory to ComfyUI
WORKDIR /comfyui

ARG SKIP_DEFAULT_MODELS
# Download checkpoints/vae/LoRA to include in image.
RUN if [ -z "$SKIP_DEFAULT_MODELS" ]; then wget -O models/checkpoints/sd_xl_base_1.0.safetensors https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors; fi
RUN if [ -z "$SKIP_DEFAULT_MODELS" ]; then wget -O models/vae/sdxl_vae.safetensors https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors; fi
RUN if [ -z "$SKIP_DEFAULT_MODELS" ]; then wget -O models/vae/sdxl-vae-fp16-fix.safetensors https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors; fi

# Install ComfyUI dependencies
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 \
&& pip3 install --no-cache-dir xformers==0.0.21 \
Expand All @@ -32,14 +38,8 @@ RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https:/
# Install runpod
RUN pip3 install runpod requests

# Download checkpoints/vae/LoRA to include in image
RUN wget -O models/checkpoints/sd_xl_base_1.0.safetensors https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors
RUN wget -O models/vae/sdxl_vae.safetensors https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors
RUN wget -O models/vae/sdxl-vae-fp16-fix.safetensors https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors

# Example for adding specific models into image
# ADD models/checkpoints/sd_xl_base_1.0.safetensors models/checkpoints/
# ADD models/vae/sdxl_vae.safetensors models/vae/
# Support for the network volume
ADD src/extra_model_paths.yaml ./

# Go back to the root
WORKDIR /
Expand Down
253 changes: 185 additions & 68 deletions README.md

Large diffs are not rendered by default.

Binary file added assets/my-endpoint-with-endpointID.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
4 changes: 3 additions & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,15 @@ version: "3.8"

services:
comfyui:
image: timpietruskyblibla/runpod-worker-comfy:latest
image: timpietruskyblibla/runpod-worker-comfy:dev
container_name: comfyui-worker
environment:
- NVIDIA_VISIBLE_DEVICES=all
- SERVE_API_LOCALLY=true
ports:
- "8000:8000"
- "8188:8188"
runtime: nvidia
volumes:
- ./data/comfyui/output:/comfyui/output
- ./data/runpod-volume:/runpod-volume
11 changes: 11 additions & 0 deletions src/extra_model_paths.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
runpod_worker_comfy:
base_path: /runpod-volume
checkpoints: models/checkpoints/
clip: models/clip/
clip_vision: models/clip_vision/
configs: models/configs/
controlnet: models/controlnet/
embeddings: models/embeddings/
loras: models/loras/
upscale_models: models/upscale_models/
vae: models/vae/
6 changes: 3 additions & 3 deletions src/rp_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@
# Maximum number of API check attempts
COMFY_API_AVAILABLE_MAX_RETRIES = 500
# Time to wait between poll attempts in milliseconds
COMFY_POLLING_INTERVAL_MS = 250
COMFY_POLLING_INTERVAL_MS = os.environ.get("COMFY_POLLING_INTERVAL_MS", 250)
# Maximum number of poll attempts
COMFY_POLLING_MAX_RETRIES = 500
COMFY_POLLING_MAX_RETRIES = os.environ.get("COMFY_POLLING_MAX_RETRIES", 500)
# Host where ComfyUI is running
COMFY_HOST = "127.0.0.1:8188"
# Enforce a clean state after each job is done
Expand Down Expand Up @@ -237,7 +237,7 @@ def process_output_images(outputs, job_id):
for node_id, node_output in outputs.items():
if "images" in node_output:
for image in node_output["images"]:
output_images = f"{image['subfolder']}/{image['filename']}"
output_images = os.path.join(image["subfolder"], image["filename"])

print(f"runpod-worker-comfy - image generation is done")

Expand Down
13 changes: 8 additions & 5 deletions src/start.sh
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,17 @@
TCMALLOC="$(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1)"
export LD_PRELOAD="${TCMALLOC}"

echo "runpod-worker-comfy: Starting ComfyUI"
python3 /comfyui/main.py --disable-auto-launch --disable-metadata &

echo "runpod-worker-comfy: Starting RunPod Handler"

# Serve the API and don't shutdown the container
if [ "$SERVE_API_LOCALLY" == "true" ]; then
echo "runpod-worker-comfy: Starting ComfyUI"
python3 /comfyui/main.py --disable-auto-launch --disable-metadata --listen &

echo "runpod-worker-comfy: Starting RunPod Handler"
python3 -u /rp_handler.py --rp_serve_api --rp_api_host=0.0.0.0
else
echo "runpod-worker-comfy: Starting ComfyUI"
python3 /comfyui/main.py --disable-auto-launch --disable-metadata &

echo "runpod-worker-comfy: Starting RunPod Handler"
python3 -u /rp_handler.py
fi
84 changes: 84 additions & 0 deletions test_resources/workflows/workflow_sdxl_turbo.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
{
"input": {
"workflow": {
"3": {
"inputs": {
"seed": 457699577674669,
"steps": 3,
"cfg": 1.5,
"sampler_name": "euler_ancestral",
"scheduler": "normal",
"denoise": 1,
"model": ["4", 0],
"positive": ["6", 0],
"negative": ["7", 0],
"latent_image": ["5", 0]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"4": {
"inputs": {
"ckpt_name": "sd_xl_turbo_1.0_fp16.safetensors"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {
"title": "Load Checkpoint"
}
},
"5": {
"inputs": {
"width": 1024,
"height": 1024,
"batch_size": 1
},
"class_type": "EmptyLatentImage",
"_meta": {
"title": "Empty Latent Image"
}
},
"6": {
"inputs": {
"text": "ancient rome, 4k photo",
"clip": ["4", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"7": {
"inputs": {
"text": "text, watermark, blurry, ugly, deformed",
"clip": ["4", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": ["3", 0],
"vae": ["4", 2]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"9": {
"inputs": {
"filename_prefix": "images/rome",
"images": ["8", 0]
},
"class_type": "SaveImage",
"_meta": {
"title": "Save Image"
}
}
}
}
}
114 changes: 114 additions & 0 deletions test_resources/workflows/workflow_webp.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
{
"input": {
"workflow": {
"3": {
"inputs": {
"seed": 416138702284529,
"steps": 20,
"cfg": 8,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 1,
"model": ["4", 0],
"positive": ["6", 0],
"negative": ["7", 0],
"latent_image": ["5", 0]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"4": {
"inputs": {
"ckpt_name": "v1-5-pruned-emaonly.safetensors"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {
"title": "Load Checkpoint"
}
},
"5": {
"inputs": {
"width": 512,
"height": 512,
"batch_size": 1
},
"class_type": "EmptyLatentImage",
"_meta": {
"title": "Empty Latent Image"
}
},
"6": {
"inputs": {
"text": "beautiful scenery nature glass bottle landscape, purple galaxy bottle,",
"clip": ["4", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"7": {
"inputs": {
"text": "text, watermark",
"clip": ["4", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": ["3", 0],
"vae": ["4", 2]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"10": {
"inputs": {
"output_path": "output_path",
"filename_prefix": "filename_prefix",
"filename_delimiter": "___",
"filename_number_padding": 4,
"filename_number_start": "false",
"extension": "webp",
"quality": 10,
"lossless_webp": "false",
"overwrite_mode": "false",
"show_history": "false",
"show_history_by_prefix": "false",
"embed_workflow": "false",
"show_previews": "false",
"images": ["8", 0]
},
"class_type": "Image Save",
"_meta": {
"title": "Image Save"
}
},
"11": {
"inputs": {
"images": ["8", 0]
},
"class_type": "PreviewImage",
"_meta": {
"title": "Preview Image"
}
},
"12": {
"inputs": {
"images": ["8", 0]
},
"class_type": "PreviewImage",
"_meta": {
"title": "Preview Image"
}
}
}
}
}
8 changes: 6 additions & 2 deletions tests/test_rp_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,9 @@ def test_bucket_endpoint_not_configured(self, mock_upload_image, mock_exists):
mock_exists.return_value = True
mock_upload_image.return_value = "simulated_uploaded/image.png"

outputs = {"node_id": {"images": [{"filename": "ComfyUI_00001_.png", "subfolder": ""}]}}
outputs = {
"node_id": {"images": [{"filename": "ComfyUI_00001_.png", "subfolder": ""}]}
}
job_id = "123"

result = rp_handler.process_output_images(outputs, job_id)
Expand Down Expand Up @@ -188,7 +190,9 @@ def test_bucket_image_upload_fails_env_vars_wrong_or_missing(
# When AWS credentials are wrong or missing, upload_image should return 'simulated_uploaded/...'
mock_upload_image.return_value = "simulated_uploaded/image.png"

outputs = {"node_id": {"images": [{"filename": "ComfyUI_00001_.png", "subfolder": "test"}]}}
outputs = {
"node_id": {"images": [{"filename": "ComfyUI_00001_.png", "subfolder": ""}]}
}
job_id = "123"

result = rp_handler.process_output_images(outputs, job_id)
Expand Down

0 comments on commit 070cde5

Please sign in to comment.