Pruna nodes for ComfyUI

This repository explains how to accelerate image generation in ComfyUI using Pruna, an inference optimization engine that makes AI models faster, smaller, cheaper, and greener. ComfyUI is a popular node-based GUI for image generation models, for which we provide a custom compilation node that accelerates Stable Diffusion (SD) and Flux inference, while preserving output quality.

Here, you'll find:

• Installation and usage instructions
• Representative workflows
• Performance benchmarks

Installation

Prerequisites

1. Create a new conda environment with Python 3.10
2. Install ComfyUI
3. Install the latest version of Pruna
4. Generate a Pruna token, following the instructions here

Steps

1. Navigate to your ComfyUI installation's custom_nodes folder:

cd <path_to_comfyui>/custom_nodes

2. Clone this repository:

git clone https://github.com/PrunaAI/ComfyUI_pruna.git

3. Launch ComfyUI, for example, with:

cd <path_to_comfyui> && python main.py --disable-cuda-malloc --gpu-only

The Pruna node will appear in the nodes menu in the Pruna category.

Important note: The current implementation requires launching ComfyUI with the --disable-cuda-malloc flag; otherwise the node may not function properly. For optimal performance, we also recommend setting the --gpu-only flag.

Usage Example

We provide two example workflows: one using a Stable Diffusion model and another based on Flux.

To load the workflow:

• Drag and drop the provided json file (for SD or Flux) into the ComfyUI window
• OR Click Open in the Workflow tab, as shown here, and select the file

Through the GUI, you can choose your preferred compilation mode. Currently, we support x-fast and torch_compile, with x-fast set as the default.

Model Setup

Example 1: Stable Diffusion

You have two options for the base model:

Option 1: SafeTensors Format (Recommended)

1. Download the safetensors version
2. Place it in <path_to_comfyui>/models/checkpoints

Option 2: Diffusers Format

1. Download the Diffusers version of SD v1.4
2. Place it in <path_to_comfyui>/models/diffusers
3. Replace the Load Checkpoint node with a DiffusersLoader node

The node is tested using the SafeTensors format, so for the sake of reproducibility, we recommend using that format. However, we don't expect any performance differences between the two.

Note: In this example, we use the Stable Diffusion v1.4 model. However, our node is compatible with any other SD model — feel free to use your favorite one!

Example 2: Flux

To use Flux, you must separately download all model components—including the VAE, CLIP, and diffusion model weights—and place them in the appropriate folder.

Steps to set up Flux:

1. For the CLIP models: Get the following files:

   • clip_l.safetensors
   • t5xxl_fp16.safetensors

   Move them to <path_to_comfyui>/models/clip/.

2. For the VAE model: Get the VAE model, and move it to <path_to_comfyui>/models/vae/ directory.

3. For the Flux model: You first need to request access to the model here. Once you have access, download the weights and move them to <path_to_comfyui>/models/diffusion_models/.

Performance

The node was tested on an NVIDIA L40S GPU. Below, we compare the performance of the base model, with the models compiled with Pruna's x-fast and torch_compile compilation modes.

• The first figure presents the performance in iterations per second, as reported by ComfyUI.
• The second figure shows the end-to-end time required to generate a single image.

Contact

For questions, feedback or community discussions, feel free to join our Discord.

For bug reports or technical issues, please open an issue in this repository.