Brev.dev, inc of San Francisco, California, USA has been acquired by NVIDIA Corporation of Santa Clara, California, USA on July 2024lawyers, are you happy now?

Learn more
DocsBlogPricing

AI

Run ComfyUI with TensorRT on Brev 🤙

Anish Maddipoti

Anish Maddipoti

June 4, 202414 min read

What is TensorRT?

TensorRT is a library, essentially a high-performance deep learning inference optimizer and runtime engine, developed by NVIDIA to maximize the performance of deep learning models on NVIDIA GPUs, enabling faster inference and lower latency for deployments.

TensorRT works by optimizing trained deep learning models from various frameworks like TensorFlow, PyTorch, and ONNX for specific hardware! It applies techniques such as layer and tensor fusion, kernel auto-tuning, precision calibration, and dynamic tensor memory to optimize the model for inference on NVIDIA GPUs. The optimized model is then deployed using the highly-parallel TensorRT inference runtime, which can deliver up to 36x faster inference compared to CPU-only platforms. TensorRT is widely adopted across industries like autonomous vehicles, robotics, healthcare, and more, where low-latency inference is critical.

Recently, ComfyUI has been integrated with TensorRT to enable faster inference and better performance for image and video generation workflows.

What is ComfyUI?

ComfyUI is a powerful user interface designed specifically for Stable Diffusion (SD), an AI model used for generating images from text prompts. It provides an intuitive workflow environment that allows users to create complex pipelines for image or video generation tasks.

ComfyUI is a visually connected interface where different components, such as checkpoint nodes, prompt inputs, and sampling methods, can be combined to form a customized workflow. Users can load SD models, input prompts, adjust parameters, and monitor GPU performance within the interface. The generated images are then displayed and can be saved for further use or analysis. ComfyUI's flexibility, optimization for GPU performance, and support for the latest Stable Diffusion features make it a popular choice among AI artists and researchers. It's become the GUI of choice for many users who want to generate high-quality images quickly and efficiently through complex workflows!

How does TensorRT work with ComfyUI?

TensorRT is being integrated into ComfyUI to accelerate the speed & performance of SD models on NVIDIA GPUs. ComfyUI now has nodes that allow users to convert their SD checkpoints into optimized TensorRT engines, with options for static engines with fixed resolution and batch size, or dynamic engines that support a range of resolutions and batch sizes. Once converted, these optimized TensorRT engines can be loaded directly into ComfyUI workflows using the TensorRT Loader node. By leveraging TensorRT's optimizations, the SD model can then take advantage of significantly accelerated performance. This integration currently supports various Stable Diffusion models, including SD 1.5, 2.1, SDXL, SDXL Turbo, Stable Video Diffusion, and Stable Video Diffusion-XT, although compatibility with ControlNets and LoRAs has been left for future updates. To benefit from TensorRT acceleration, users need an NVIDIA GeForce RTX or NVIDIA RTX GPU, with higher VRAM GPUs (12GB+) recommended for larger models like SDXL. On Brev, we have a ton of RTX GPUs available, ranging from 4090s all the way up to 8000s that meet the VRAM requurements!

How to run ComfyUI with TensorRT on Brev?

Enough talk, let's get to the action! To run ComfyUI with TensorRT on Brev, follow these steps:

Click on this Launchable to launch a RTX GPU on Brev with a ComfyUI x TensorRT Jupyter Notebook pre-loaded! I'll still put in some of the code from the notebook in this blog post and you can also find the full notebook here.

Let's start by installing the necessary dependencies.

# Environment Setup

from pathlib import Path

OPTIONS = {}

WORKSPACE = 'ComfyUI'
USE_COMFYUI_MANAGER = True  #@param {type:"boolean"}
UPDATE_COMFY_UI = True  #@param {type:"boolean"}
INSTALL_CUSTOM_NODES_DEPENDENCIES = True  #@param {type:"boolean"}
OPTIONS['UPDATE_COMFY_UI'] = UPDATE_COMFY_UI
OPTIONS['USE_COMFYUI_MANAGER'] = USE_COMFYUI_MANAGER
OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES'] = INSTALL_CUSTOM_NODES_DEPENDENCIES

Let's clone the ComfyUI repository and install the necessary dependencies.

%%bash
if ! [ -d WORKSPACE ]; then
    echo "Starting initial setup of ComfyUI..."
    git clone https://github.com/comfyanonymous/ComfyUI
    if [ -d ComfyUI ]; then
        echo "Repository cloned successfully into 'ComfyUI'."
        cd ComfyUI
    else
        echo "Failed to clone repository."
    fi
else
    echo "Workspace directory already exists."
fi

if OPTIONS['UPDATE_COMFY_UI']:
    !echo "Updating ComfyUI..."
    !git pull

!echo "Installing dependencies..."
!pip install accelerate
!pip install einops
!pip install transformers
!pip install safetensors>=0.3.0
!pip install aiohttp
!pip install pyyaml
!pip install Pillow
!pip install scipy
!pip install tqdm
!pip install psutil
!pip install torch==2.2.2 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
!pip install torchsde
!pip install kornia>=0.7.1 spandrel
!pip install typer
!pip install diffusers

Optionally, you can also download ComfyUI manager, which will allow you to easily manage your ComfyUI installation and custom nodes.

import os

if not os.path.isfile("ComfyUI-Manager/check.sh"):
    !chmod 755 ComfyUI-Manager/check.sh
if not os.path.isfile("ComfyUI-Manager/scan.sh"):
    !chmod 755 ComfyUI-Manager/scan.sh
if not os.path.isfile("ComfyUI-Manager/node_db/dev/scan.sh"):
    !chmod 755 ComfyUI-Manager/node_db/dev/scan.sh
if not os.path.isfile("ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh"):
    !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh
if not os.path.isfile("ComfyUI-Manager/scripts/install-comfyui-venv-win.bat"):
    !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-win.bat

if not os.path.isdir("ComfyUI-Manager"):
    print("-= Initial setup ComfyUI-Manager =-")
    !git clone https://github.com/ltdrdata/ComfyUI-Manager

%cd ComfyUI-Manager
!git pull

%cd ..

%cd ~/verb-workspace

import shutil
import os

# Move the ComfyUI-Manager folder to the custom_nodes directory
source_manager_folder = '/root/verb-workspace/ComfyUI-Manager'
# source_manager_folder = '/home/ubuntu/verb-workspace/ComfyUI-Manager' (if you run into error, try this)
destination_manager_folder = '/root/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager'
# destination_manager_folder = '/home/ubuntu/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager' (if you run into error, try this)

if os.path.exists(destination_manager_folder):
    shutil.rmtree(destination_manager_folder)

shutil.move(source_manager_folder, destination_manager_folder)

assert os.path.exists(destination_manager_folder), "Folder move failed for ComfyUI-Manager"
print(f'Folder ComfyUI-Manager moved to {destination_manager_folder}')

if 'INSTALL_CUSTOM_NODES_DEPENDENCIES' in OPTIONS and OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES']:
    print("-= Install custom nodes dependencies =-")
    !pip install GitPython
    !python /root/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager/cm-cli.py restore-dependencies

Let's download some checkpoints from HuggingFace!

!pip install huggingface_hub

!wget -c https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -P ./ComfyUI/models/checkpoints/

!wget -c https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors -P ./ComfyUI/models/vae/

We're going to build the TensorRT engine for this next model below, Stable Diffusion XL Turbo.

from diffusers import StableDiffusionPipeline

# Download the model
!wget -O ./ComfyUI/models/checkpoints/sd_xl_turbo_1.0_fp16.safetensors https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/sd_xl_turbo_1.0_fp16.safetensors

# Load the model from the local path
pipeline = StableDiffusionPipeline.from_single_file(
    "./ComfyUI/models/checkpoints/sd_xl_turbo_1.0_fp16.safetensors"
)

!cd ./ComfyUI/custom_nodes && git clone https://github.com/comfyanonymous/ComfyUI_TensorRT

!cd ./ComfyUI/custom_nodes/ComfyUI_TensorRT && pip install -r requirements.txt

We need to use 2 JSON files to build the TensorRT engine and actually run it in a ComfyUI workflow.

Let's pull them from the TensorRT repository and move them to the main directory, so that you can download them locally. You'll need to load these JSON files into ComfyUI (which I'll show you how to do in the next section).

source_folder_path = '/root/verb-workspace/ComfyUI/custom_nodes/ComfyUI_TensorRT/workflows'
# source_folder_path = '/home/ubuntu/verb-workspace/ComfyUI/custom_nodes/ComfyUI_TensorRT/workflows' (if you run into error, try this)
destination_folder_path = '/root/verb-workspace/'
# destination_folder_path = '/home/ubuntu/verb-workspace' (if you run into error, try this)

json_files = [
    'Create_SDXL_Turbo_TRT_Static.json',
    'Build.TRT.Engine_SDXL_Turbo_Static.json'
]

# Move the workflows folder to the main directory
destination_workflows_path = os.path.join(destination_folder_path, 'workflows')
if os.path.exists(destination_workflows_path):
    shutil.rmtree(destination_workflows_path)

shutil.move(source_folder_path, destination_folder_path)

# Move the JSON files out of the workflows folder into the main directory
for json_file in json_files:
    source_json_path = os.path.join(destination_workflows_path, json_file)
    destination_json_path = os.path.join(destination_folder_path, json_file)
    shutil.move(source_json_path, destination_json_path)

    assert os.path.exists(destination_json_path), f"File move failed for {json_file}"
    print(f'File {json_file} moved to {destination_json_path}')

Let's start the ComfyUI server!

import subprocess
import threading
import time
import socket

def check_server(port):
    while True:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            result = sock.connect_ex(('127.0.0.1', port))
            if result == 0:
                print("\nServer is up and running at http://129.80.69.99:{}".format(port))
                break
            time.sleep(0.5)

threading.Thread(target=check_server, daemon=True, args=(8188,)).start()

This will start the server on port 8188.

!python /root/verb-workspace/ComfyUI/main.py --listen 0.0.0.0 --port 8188 --dont-print-server --disable-xformers
Now click on the URL on your Brev instance, that's mapped to port 8188. You'll now see the ComfyUI interface with the default workflow!
img

which should take you to this:

img
You need to click on the Load button and then select the Build.TRT.Engine_SDXL_Turbo_Static.json file that you downloaded earlier.

You should see this:

img
Time to build the TensorRT engine for SDXL Turbo! Click on the Queue Prompt button and then wait a couple of minutes for the engine to be built. When it's done, a green outline will appear around the STATIC_TRT_MODEL_CONVERSION node.
img

After the engine is built (takes ~5 minutes), you need to load in the Create.TRT.Model_SDXL_Turbo_Static.json file that you downloaded earlier. Click on the Load button and then select the file from your computer. This will load the workflow with the TensorRT engine node so you can actually use it!

img

This is what my workflow looks like and the prompt I added! Click Queue Prompt to run it!

That's it! You've just run ComfyUI with TensorRT on Brev!

You can now take advantage of blazing fast inference and better performance for your ComfyUI workflows. We'll have many more resources and guides on how to use TensorRT with ComfyUI in the future, so stay tuned! If you have any questions, reach out to us on Discord. 🤙

Previous
Convert fine-tuned LLMs into GGUF to run locally with Ollama
Next
How to fine-tune Llama3 using DPO on Brev 🤙