feat/restore-snapshot (#30)

* 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 <[email protected]> * 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" * Implement optional restoring of ComfyUI snapshots * feat: use comfy-cli to install ComfyUI & restore snapshot * fix: install all dependencies from the snapshot * chore: moved example snapshot to test_resources * feat: allow any kind of snapshot file * feat: allow any file that has "snapshot" in its name * ci: added test for "restore-snapshot" * ci: run restore snapshat test automatically * docs: use "snapshots" to bake custom nodes into the docker image * docs: add link to ComfyUI Manager docs on how to export snapshots * ci: use correct path to example_snapshot.json * ci: fix the path * ci: fix the file path * ci: use a mock instead of the actual file * ci: fix the path * chore: don't ignore snapshot.json as people might use this --------- Co-authored-by: Tim Pietrusky <[email protected]> Co-authored-by: Meptl <[email protected]>
blib-la · Nov 18, 2024 · 7e4b7d0 · 7e4b7d0
1 parent 06164a3
commit 7e4b7d0
Show file tree

Hide file tree

Showing 6 changed files with 402 additions and 58 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -23,5 +23,10 @@ jobs:
       - name: Install dependencies
         run: pip install -r requirements.txt
 
-      - name: Run tests
+      - name: Run Python tests
         run: python -m unittest discover
+
+      - name: Run snapshot restoration tests
+        run: |
+          chmod +x tests/test_restore_snapshot.sh
+          ./tests/test_restore_snapshot.sh
diff --git a/Dockerfile b/Dockerfile
@@ -7,39 +7,49 @@ ENV DEBIAN_FRONTEND=noninteractive
 ENV PIP_PREFER_BINARY=1
 # Ensures output from python is printed immediately to the terminal without buffering
 ENV PYTHONUNBUFFERED=1 
+# Speed up some cmake builds
+ENV CMAKE_BUILD_PARALLEL_LEVEL=8
 
 # Install Python, git and other necessary tools
 RUN apt-get update && apt-get install -y \
     python3.10 \
     python3-pip \
     git \
-    wget
+    wget \
+    libgl1 \
+    && ln -sf /usr/bin/python3.10 /usr/bin/python \
+    && ln -sf /usr/bin/pip3 /usr/bin/pip
 
 # Clean up to reduce image size
 RUN apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/*
 
-# Clone ComfyUI repository
-RUN git clone https://github.com/comfyanonymous/ComfyUI.git /comfyui
+# Install comfy-cli
+RUN pip install comfy-cli
+
+# Install ComfyUI
+RUN /usr/bin/yes | comfy --workspace /comfyui install --cuda-version 11.8 --nvidia --version 0.2.7
 
 # Change working directory to ComfyUI
 WORKDIR /comfyui
 
-# Install ComfyUI dependencies
-RUN pip3 install --upgrade --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 \
-    && pip3 install --upgrade -r requirements.txt
-
 # Install runpod
-RUN pip3 install runpod requests
+RUN pip install runpod requests
 
 # Support for the network volume
 ADD src/extra_model_paths.yaml ./
 
 # Go back to the root
 WORKDIR /
 
-# Add the start and the handler
-ADD src/start.sh src/rp_handler.py test_input.json ./
-RUN chmod +x /start.sh
+# Add scripts
+ADD src/start.sh src/restore_snapshot.sh src/rp_handler.py test_input.json ./
+RUN chmod +x /start.sh /restore_snapshot.sh
+
+# Optionally copy the snapshot file
+ADD *snapshot*.json /
+
+# Restore the snapshot to install custom nodes
+RUN /restore_snapshot.sh
 
 # Stage 2: Download models
 FROM base as downloader
@@ -79,4 +89,3 @@ FROM base as final
 COPY --from=downloader /comfyui/models /comfyui/models
 
 # Start the container
-CMD /start.sh
diff --git a/README.md b/README.md
@@ -10,39 +10,40 @@ Read our article here: https://blib.la/blog/comfyui-on-runpod
 
 [![Discord](https://img.shields.io/discord/1091306623819059300?color=7289da&label=Discord&logo=discord&logoColor=fff&style=for-the-badge)](https://discord.com/invite/m3TBB9XEkb)
 
-→ Please also checkout [Captain: The AI Platform](https://github.com/blib-la/captain)
-
 ---
 
 <!-- toc -->
 
 - [Quickstart](#quickstart)
 - [Features](#features)
 - [Config](#config)
-  * [Upload image to AWS S3](#upload-image-to-aws-s3)
+  - [Upload image to AWS S3](#upload-image-to-aws-s3)
 - [Use the Docker image on RunPod](#use-the-docker-image-on-runpod)
-  * [Create your template (optional)](#create-your-template-optional)
-  * [Create your endpoint](#create-your-endpoint)
-  * [GPU recommendations](#gpu-recommendations)
+  - [Create your template (optional)](#create-your-template-optional)
+  - [Create your endpoint](#create-your-endpoint)
+  - [GPU recommendations](#gpu-recommendations)
 - [API specification](#api-specification)
-  * [JSON Request Body](#json-request-body)
-  * [Fields](#fields)
-    + ["input.images"](#inputimages)
+  - [JSON Request Body](#json-request-body)
+  - [Fields](#fields)
+    - ["input.images"](#inputimages)
 - [Interact with your RunPod API](#interact-with-your-runpod-api)
-  * [Health status](#health-status)
-  * [Generate an image](#generate-an-image)
-    + [Example request for SDXL with cURL](#example-request-for-sdxl-with-curl)
+  - [Health status](#health-status)
+  - [Generate an image](#generate-an-image)
+    - [Example request for SDXL with cURL](#example-request-for-sdxl-with-curl)
 - [How to get the workflow from ComfyUI?](#how-to-get-the-workflow-from-comfyui)
 - [Bring Your Own Models and Nodes](#bring-your-own-models-and-nodes)
-  * [Network Volume](#network-volume)
-  * [Custom Docker Image](#custom-docker-image)
+  - [Network Volume](#network-volume)
+  - [Custom Docker Image](#custom-docker-image)
+    - [Adding Custom Models](#adding-custom-models)
+    - [Adding Custom Nodes](#adding-custom-nodes)
+    - [Building the Image](#building-the-image)
 - [Local testing](#local-testing)
-  * [Setup](#setup)
-    + [Setup for Windows](#setup-for-windows)
-  * [Testing the RunPod handler](#testing-the-runpod-handler)
-  * [Local API](#local-api)
-    + [Access the local Worker API](#access-the-local-worker-api)
-    + [Access local ComfyUI](#access-local-comfyui)
+  - [Setup](#setup)
+    - [Setup for Windows](#setup-for-windows)
+  - [Testing the RunPod handler](#testing-the-runpod-handler)
+  - [Local API](#local-api)
+    - [Access the local Worker API](#access-the-local-worker-api)
+    - [Access local ComfyUI](#access-local-comfyui)
 - [Automatically deploy to Docker hub with GitHub Actions](#automatically-deploy-to-docker-hub-with-github-actions)
 - [Acknowledgments](#acknowledgments)
 
@@ -293,39 +294,52 @@ Note: The folders in the Network Volume are automatically available to ComfyUI w
 
 ### Custom Docker Image
 
-If you prefer to include your models directly in the Docker image, follow these steps:
+If you prefer to include your models and custom nodes directly in the Docker image, follow these steps:
 
 1. **Fork the Repository**:
-
    - Fork this repository to your own GitHub account.
 
-2. **Add Your Models in the Dockerfile**:
+#### Adding Custom Models
 
-   - Edit the `Dockerfile` to include your models:
-     ```Dockerfile
-     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
-     ```
-   - You can also add custom nodes:
-     ```Dockerfile
-     RUN git clone https://github.com/<username>/<custom-node-repo>.git custom_nodes/<custom-node-repo>
-     ```
+To include additional models in your Docker image, edit the `Dockerfile` and add the download commands:
 
-3. **Build Your Docker Image**:
-   - Build the **base** image locally:
-     ```bash
-     docker build -t <your_dockerhub_username>/runpod-worker-comfy:dev-base --target base --platform linux/amd64 .
-     ```
-   - Build the **sdxl** image locally:
-     ```bash
-     docker build --build-arg MODEL_TYPE=sdxl -t <your_dockerhub_username>/runpod-worker-comfy:dev-sdxl --platform linux/amd64 .
-     ```
-   - Build the **sd3** image locally:
-     ```bash
-     docker build --build-arg MODEL_TYPE=sd3 --build-arg HUGGINGFACE_ACCESS_TOKEN=<your-huggingface-token> -t <your_dockerhub_username>/runpod-worker-comfy:dev-sd3 --platform linux/amd64 .
-     ```
+```Dockerfile
+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
+```
+
+#### Adding Custom Nodes
+
+To include custom nodes in your Docker image:
+
+1. [Export a snapshot from ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager?tab=readme-ov-file#snapshot-manager) that includes all your desired custom nodes
+   a. Open "Manager > Snapshot Manager"
+   b. Create a new snapshot by clicking on "Save snapshot"
+   c. Get the `*_snapshot.json` from your ComfyUI: `ComfyUI/custom_nodes/ComfyUI-Manager/snapshots`
+2. Save the snapshot file in the root directory of the project
+3. The snapshot will be automatically restored during the Docker build process, see [Building the Image](#building-the-image)
+
+> [!NOTE]
+>
+> - Some custom nodes may download additional models during installation, which can significantly increase the image size
+> - Having many custom nodes may increase ComfyUI's initialization time
+
+#### Building the Image
+
+Build your customized Docker image locally:
+
+```bash
+# Build the base image
+docker build -t <your_dockerhub_username>/runpod-worker-comfy:dev-base --target base --platform linux/amd64 .
+
+# Build the SDXL image
+docker build --build-arg MODEL_TYPE=sdxl -t <your_dockerhub_username>/runpod-worker-comfy:dev-sdxl --platform linux/amd64 .
+
+# Build the SD3 image
+docker build --build-arg MODEL_TYPE=sd3 --build-arg HUGGINGFACE_ACCESS_TOKEN=<your-huggingface-token> -t <your_dockerhub_username>/runpod-worker-comfy:dev-sd3 --platform linux/amd64 .
+```
 
 > [!NOTE]  
-> Ensure to specify `--platform linux/amd64` to avoid errors on RunPod, see [issue #13](https://github.com/blib-la/runpod-worker-comfy/issues/13).
+> Ensure to specify `--platform linux/amd64` to avoid errors on RunPod, see [issue #13](https://github.com/blib-la/runpod-worker-comfy/issues/13)
 
 ## Local testing
 

diff --git a/src/restore_snapshot.sh b/src/restore_snapshot.sh
@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+set -e
+
+SNAPSHOT_FILE=$(ls /*snapshot*.json 2>/dev/null | head -n 1)
+
+if [ -z "$SNAPSHOT_FILE" ]; then
+    echo "runpod-worker-comfy: No snapshot file found. Exiting..."
+    exit 0
+fi
+
+echo "runpod-worker-comfy: restoring snapshot: $SNAPSHOT_FILE"
+
+comfy --workspace /comfyui node restore-snapshot "$SNAPSHOT_FILE" --pip-non-url
+
+echo "runpod-worker-comfy: restored snapshot file: $SNAPSHOT_FILE"