diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8fef7af..cba3c08 100644
--- 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
index de6befa..b005655 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -7,29 +7,33 @@ 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 ./
@@ -37,9 +41,15 @@ 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
\ No newline at end of file
diff --git a/README.md b/README.md
index e539ff6..0a15146 100644
--- a/README.md
+++ b/README.md
@@ -10,8 +10,6 @@ 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 -->
@@ -19,30 +17,33 @@ Read our article here: https://blib.la/blog/comfyui-on-runpod
 - [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
new file mode 100644
index 0000000..73ba068
--- /dev/null
+++ 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"
\ No newline at end of file
diff --git a/test_resources/example_snapshot.json b/test_resources/example_snapshot.json
new file mode 100644
index 0000000..24d9a21
--- /dev/null
+++ b/test_resources/example_snapshot.json
@@ -0,0 +1,238 @@
+{
+  "comfyui": "4ac401af2b3962f7a669a6757a804f7773833c47",
+  "git_custom_nodes": {
+    "https://github.com/ltdrdata/ComfyUI-Manager": {
+      "hash": "b6a8e6ba8147080a320b1b91c93a0b1cbdb93136",
+      "disabled": false
+    },
+    "https://github.com/ltdrdata/ComfyUI-Impact-Pack": {
+      "hash": "24de5a846bb9b1c1830af0218f76266c375cc904",
+      "disabled": false
+    },
+    "https://github.com/ltdrdata/ComfyUI-Inspire-Pack": {
+      "hash": "3bcc1d43a033b3d4b3066c34c9fba2a952b54ad0",
+      "disabled": false
+    }
+  },
+  "file_custom_nodes": [
+    {
+      "filename": "websocket_image_save.py",
+      "disabled": false
+    }
+  ],
+  "pips": {
+    "aiohappyeyeballs==2.4.3": "",
+    "aiohttp==3.11.2": "",
+    "aiosignal==1.3.1": "",
+    "anyio==4.1.0": "",
+    "argon2-cffi==23.1.0": "",
+    "argon2-cffi-bindings==21.2.0": "",
+    "arrow==1.3.0": "",
+    "asttokens==2.4.1": "",
+    "async-lru==2.0.4": "",
+    "async-timeout==5.0.1": "",
+    "attrs==23.1.0": "",
+    "Babel==2.13.1": "",
+    "beautifulsoup4==4.12.2": "",
+    "bleach==6.1.0": "",
+    "blinker==1.4": "",
+    "cachetools==5.5.0": "",
+    "certifi==2023.11.17": "",
+    "cffi==1.16.0": "",
+    "charset-normalizer==3.3.2": "",
+    "click==8.1.7": "",
+    "comfy-cli==1.3.1": "",
+    "comm==0.2.0": "",
+    "contourpy==1.3.1": "",
+    "cryptography==3.4.8": "",
+    "cycler==0.12.1": "",
+    "dbus-python==1.2.18": "",
+    "debugpy==1.8.0": "",
+    "decorator==5.1.1": "",
+    "defusedxml==0.7.1": "",
+    "Deprecated==1.2.15": "",
+    "dill==0.3.9": "",
+    "distro==1.7.0": "",
+    "einops==0.8.0": "",
+    "entrypoints==0.4": "",
+    "exceptiongroup==1.2.0": "",
+    "executing==2.0.1": "",
+    "fastjsonschema==2.19.0": "",
+    "filelock==3.13.1": "",
+    "fonttools==4.55.0": "",
+    "fqdn==1.5.1": "",
+    "frozenlist==1.5.0": "",
+    "fsspec==2023.10.0": "",
+    "gitdb==4.0.11": "",
+    "GitPython==3.1.43": "",
+    "h11==0.14.0": "",
+    "httpcore==1.0.7": "",
+    "httplib2==0.20.2": "",
+    "httpx==0.27.2": "",
+    "huggingface-hub==0.26.2": "",
+    "idna==3.6": "",
+    "imageio==2.36.0": "",
+    "importlib-metadata==4.6.4": "",
+    "ipykernel==6.26.0": "",
+    "ipython==8.18.1": "",
+    "ipython-genutils==0.2.0": "",
+    "ipywidgets==8.1.1": "",
+    "isoduration==20.11.0": "",
+    "jedi==0.19.1": "",
+    "jeepney==0.7.1": "",
+    "Jinja2==3.1.2": "",
+    "json5==0.9.14": "",
+    "jsonpointer==2.4": "",
+    "jsonschema==4.20.0": "",
+    "jsonschema-specifications==2023.11.1": "",
+    "jupyter-archive==3.4.0": "",
+    "jupyter-contrib-core==0.4.2": "",
+    "jupyter-contrib-nbextensions==0.7.0": "",
+    "jupyter-events==0.9.0": "",
+    "jupyter-highlight-selected-word==0.2.0": "",
+    "jupyter-lsp==2.2.1": "",
+    "jupyter-nbextensions-configurator==0.6.3": "",
+    "jupyter_client==7.4.9": "",
+    "jupyter_core==5.5.0": "",
+    "jupyter_server==2.10.1": "",
+    "jupyter_server_terminals==0.4.4": "",
+    "jupyterlab==4.0.9": "",
+    "jupyterlab-widgets==3.0.9": "",
+    "jupyterlab_pygments==0.3.0": "",
+    "jupyterlab_server==2.25.2": "",
+    "keyring==23.5.0": "",
+    "kiwisolver==1.4.7": "",
+    "kornia==0.7.4": "",
+    "kornia_rs==0.1.7": "",
+    "launchpadlib==1.10.16": "",
+    "lazr.restfulclient==0.14.4": "",
+    "lazr.uri==1.0.6": "",
+    "lazy_loader==0.4": "",
+    "lxml==4.9.3": "",
+    "markdown-it-py==3.0.0": "",
+    "MarkupSafe==2.1.3": "",
+    "matplotlib==3.9.2": "",
+    "matplotlib-inline==0.1.6": "",
+    "matrix-client==0.4.0": "",
+    "mdurl==0.1.2": "",
+    "mistune==3.0.2": "",
+    "mixpanel==4.10.1": "",
+    "more-itertools==8.10.0": "",
+    "mpmath==1.3.0": "",
+    "multidict==6.1.0": "",
+    "nbclassic==1.0.0": "",
+    "nbclient==0.9.0": "",
+    "nbconvert==7.11.0": "",
+    "nbformat==5.9.2": "",
+    "nest-asyncio==1.5.8": "",
+    "networkx==3.2.1": "",
+    "notebook==6.5.5": "",
+    "notebook_shim==0.2.3": "",
+    "numpy==1.26.2": "",
+    "nvidia-cublas-cu12==12.1.3.1": "",
+    "nvidia-cuda-cupti-cu12==12.1.105": "",
+    "nvidia-cuda-nvrtc-cu12==12.1.105": "",
+    "nvidia-cuda-runtime-cu12==12.1.105": "",
+    "nvidia-cudnn-cu12==8.9.2.26": "",
+    "nvidia-cufft-cu12==11.0.2.54": "",
+    "nvidia-curand-cu12==10.3.2.106": "",
+    "nvidia-cusolver-cu12==11.4.5.107": "",
+    "nvidia-cusparse-cu12==12.1.0.106": "",
+    "nvidia-nccl-cu12==2.18.1": "",
+    "nvidia-nvjitlink-cu12==12.3.101": "",
+    "nvidia-nvtx-cu12==12.1.105": "",
+    "oauthlib==3.2.0": "",
+    "opencv-python==4.10.0.84": "",
+    "opencv-python-headless==4.10.0.84": "",
+    "overrides==7.4.0": "",
+    "packaging==23.2": "",
+    "pandas==2.2.3": "",
+    "pandocfilters==1.5.0": "",
+    "parso==0.8.3": "",
+    "pathspec==0.12.1": "",
+    "pexpect==4.9.0": "",
+    "piexif==1.1.3": "",
+    "Pillow==10.1.0": "",
+    "platformdirs==4.0.0": "",
+    "prometheus-client==0.19.0": "",
+    "prompt-toolkit==3.0.36": "",
+    "propcache==0.2.0": "",
+    "psutil==5.9.6": "",
+    "ptyprocess==0.7.0": "",
+    "pure-eval==0.2.2": "",
+    "py-cpuinfo==9.0.0": "",
+    "pycparser==2.21": "",
+    "PyGithub==2.5.0": "",
+    "Pygments==2.17.2": "",
+    "PyGObject==3.42.1": "",
+    "PyJWT==2.9.0": "",
+    "PyNaCl==1.5.0": "",
+    "pyparsing==2.4.7": "",
+    "python-apt==2.4.0+ubuntu2": "",
+    "python-dateutil==2.8.2": "",
+    "python-json-logger==2.0.7": "",
+    "pytz==2024.2": "",
+    "PyYAML==6.0.1": "",
+    "pyzmq==24.0.1": "",
+    "questionary==2.0.1": "",
+    "referencing==0.31.0": "",
+    "regex==2024.11.6": "",
+    "requests==2.31.0": "",
+    "rfc3339-validator==0.1.4": "",
+    "rfc3986-validator==0.1.1": "",
+    "rich==13.9.4": "",
+    "rpds-py==0.13.1": "",
+    "safetensors==0.4.5": "",
+    "scikit-image==0.24.0": "",
+    "scipy==1.14.1": "",
+    "seaborn==0.13.2": "",
+    "SecretStorage==3.3.1": "",
+    "segment-anything==1.0": "",
+    "semver==3.0.2": "",
+    "Send2Trash==1.8.2": "",
+    "sentencepiece==0.2.0": "",
+    "shellingham==1.5.4": "",
+    "six==1.16.0": "",
+    "smmap==5.0.1": "",
+    "sniffio==1.3.0": "",
+    "soundfile==0.12.1": "",
+    "soupsieve==2.5": "",
+    "spandrel==0.4.0": "",
+    "stack-data==0.6.3": "",
+    "sympy==1.12": "",
+    "terminado==0.18.0": "",
+    "tifffile==2024.9.20": "",
+    "tinycss2==1.2.1": "",
+    "tokenizers==0.20.3": "",
+    "tomli==2.0.1": "",
+    "tomlkit==0.13.2": "",
+    "torch==2.1.1": "",
+    "torchaudio==2.1.1": "",
+    "torchsde==0.2.6": "",
+    "torchvision==0.16.1": "",
+    "tornado==6.3.3": "",
+    "tqdm==4.67.0": "",
+    "traitlets==5.13.0": "",
+    "trampoline==0.1.2": "",
+    "transformers==4.46.2": "",
+    "triton==2.1.0": "",
+    "typer==0.13.0": "",
+    "types-python-dateutil==2.8.19.14": "",
+    "typing_extensions==4.8.0": "",
+    "tzdata==2024.2": "",
+    "ultralytics==8.3.31": "",
+    "ultralytics-thop==2.0.11": "",
+    "uri-template==1.3.0": "",
+    "urllib3==1.26.20": "",
+    "uv==0.5.2": "",
+    "wadllib==1.3.6": "",
+    "wcwidth==0.2.12": "",
+    "webcolors==1.13": "",
+    "webencodings==0.5.1": "",
+    "websocket-client==1.6.4": "",
+    "widgetsnbextension==4.0.9": "",
+    "wrapt==1.16.0": "",
+    "yarl==1.17.1": "",
+    "zipp==1.0.0": ""
+  }
+}
diff --git a/tests/test_restore_snapshot.sh b/tests/test_restore_snapshot.sh
new file mode 100644
index 0000000..7d3b906
--- /dev/null
+++ b/tests/test_restore_snapshot.sh
@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+
+# Set up error handling
+set -e
+
+# Store the path to the script we want to test
+SCRIPT_TO_TEST="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)/src/restore_snapshot.sh"
+
+# Ensure the script exists and is executable
+if [ ! -f "$SCRIPT_TO_TEST" ]; then
+    echo "Error: Script not found at $SCRIPT_TO_TEST"
+    exit 1
+fi
+chmod +x "$SCRIPT_TO_TEST"
+
+# Create test directory
+TEST_DIR=$(mktemp -d)
+cd "$TEST_DIR"
+
+# Create a minimal mock snapshot file
+cat > test_restore_snapshot_temporary.json << 'EOF'
+{
+  "comfyui": "test-hash",
+  "git_custom_nodes": {},
+  "file_custom_nodes": [],
+  "pips": {}
+}
+EOF
+
+# Create a mock comfy command that simulates the real comfy behavior
+cat > comfy << 'EOF'
+#!/bin/bash
+if [[ "$1" == "--workspace" && "$3" == "restore-snapshot" ]]; then
+    # Verify the snapshot file exists
+    if [[ ! -f "$4" ]]; then
+        echo "Error: Snapshot file not found"
+        exit 1
+    fi
+    echo "Mock: Restored snapshot from $4"
+    exit 0
+fi
+EOF
+
+chmod +x comfy
+export PATH="$TEST_DIR:$PATH"
+
+# Run the actual restore_snapshot script
+echo "Testing snapshot restoration..."
+echo "Script location: $SCRIPT_TO_TEST"
+"$SCRIPT_TO_TEST"
+
+# Verify the script executed successfully
+if [ $? -eq 0 ]; then
+    echo "✅ Test passed: Snapshot restoration script executed successfully"
+else
+    echo "❌ Test failed: Snapshot restoration script failed"
+    exit 1
+fi
+
+# Clean up
+rm -rf "$TEST_DIR"
+  
\ No newline at end of file