feat: network-volume; execution time config; skip default images; access ComfyUI via web

Dockerfile

@@ -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 \
@@ -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 /

docker-compose.yml

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

src/extra_model_paths.yaml

@@ -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/

src/rp_handler.py

@@ -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
@@ -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")
 

src/start.sh

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

test_resources/workflows/workflow_sdxl_turbo.json

@@ -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"
+        }
+      }
+    }
+  }
+}

test_resources/workflows/workflow_webp.json

@@ -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"
+        }
+      }
+    }
+  }
+}

tests/test_rp_handler.py

@@ -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)
@@ -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)