Deploying custom ComfyUI workflows as APIs
TL;DR
We released a new feature that enables building custom ComfyUI workflows using any node or model checkpoint! You could already serve your ComfyUI model behind an API endpoint on Baseten—now, serving custom image generation pipelines is even easier.
ComfyUI is a popular GUI used to power Stable Diffusion workflows. Instead of the complexity of needing to know the diffuser’s code, ComfyUI gives you a simple user interface to run Stable Diffusion.
For most workflows using ComfyUI, the ability to run custom nodes has become essential. That’s why we built our new build_commands
feature: you can now easily package your ComfyUI workflow to use any custom node or model checkpoint on Baseten!Â
Let’s walk through an example together.
How to deploy a custom ComfyUI workflow
Style transfer is all the rage these days, so let’s deploy a style transfer workflow on Baseten that converts a picture of a pet into an anime style. The workflow we will be using can be found here.Â
First, let’s grab the ComfyUI Truss from the Truss examples Github repository:Â
git clone https://github.com/basetenlabs/truss-examples.git
cd truss-examples/comfyui-truss
This repository already contains all the files we need to deploy our ComfyUI workflow. There are just two files we need to modify: config.yaml
and data/comfy_ui_workflow.json
. Let’s start with the config.yaml
.
Step 1: Adding the build_commands
inside the config.yaml
Inside the config.yaml
file, we can specify a key called build_commands
, which are shell commands that will run during our container build process. Here is an example:
1build_commands:
2- git clone https://github.com/comfyanonymous/ComfyUI.git
3- cd ComfyUI && git checkout b1fd26fe9e55163f780bf9e5f56bf9bf5f035c93 && pip install -r requirements.txt
4- cd ComfyUI/custom_nodes && git clone https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes --recursive && cd ComfyUI-Inference-Core-Nodes && pip install -e .[cuda12]
5- cd ComfyUI/custom_nodes && git clone https://github.com/ZHO-ZHO-ZHO/ComfyUI-Gemini --recursive && cd ComfyUI-Gemini && pip install -r requirements.txt
6- cd ComfyUI/custom_nodes && git clone https://github.com/kijai/ComfyUI-Marigold --recursive && cd ComfyUI-Marigold && pip install -r requirements.txt
7- cd ComfyUI/custom_nodes && git clone https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92 --recursive
8- cd ComfyUI/custom_nodes && git clone https://github.com/Fannovel16/comfyui_controlnet_aux --recursive && cd comfyui_controlnet_aux && pip install -r requirements.txt
9- cd ComfyUI/models/controlnet && wget -O control-lora-canny-rank256.safetensors https://huggingface.co/stabilityai/control-lora/resolve/main/control-LoRAs-rank256/control-lora-canny-rank256.safetensors
10- cd ComfyUI/models/controlnet && wget -O control-lora-depth-rank256.safetensors https://huggingface.co/stabilityai/control-lora/resolve/main/control-LoRAs-rank256/control-lora-depth-rank256.safetensors
11- cd ComfyUI/models/checkpoints && wget -O dreamshaperXL_v21TurboDPMSDE.safetensors https://civitai.com/api/download/models/351306
12- cd ComfyUI/models/loras && wget -O StudioGhibli.Redmond-StdGBRRedmAF-StudioGhibli.safetensors https://huggingface.co/artificialguybr/StudioGhibli.Redmond-V2/resolve/main/StudioGhibli.Redmond-StdGBRRedmAF-StudioGhibli.safetensors
The ComfyUI manager allows you to easily install custom nodes. Behind the scenes, the manager simply clones a repository and installs the Python dependencies. We can simulate that behavior using a command like git clone ... && pip install -r requirements
. Under build_commands
, you can run arbitrary shell commands such as git clone
, cd
, or wget
. This way, you can install any checkpoints, LoRAs, and ControlNets, and place them in the appropriate folders inside ComfyUI. You can even create new directories, such as ipadapter
.
Each line under build_commands
performs a Docker RUN
command; since these are run during the image build step, they get cached into the Docker image. This way, when your container spins up, all custom nodes and models are already downloaded which significantly reduces the time for cold starts.
Here is the full config.yaml
file we’ll use for our example:
1build_commands:
2- git clone https://github.com/comfyanonymous/ComfyUI.git
3- cd ComfyUI && git checkout b1fd26fe9e55163f780bf9e5f56bf9bf5f035c93 && pip install -r requirements.txt
4- cd ComfyUI/custom_nodes && git clone https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes --recursive && cd ComfyUI-Inference-Core-Nodes && pip install -e .[cuda12]
5- cd ComfyUI/custom_nodes && git clone https://github.com/ZHO-ZHO-ZHO/ComfyUI-Gemini --recursive && cd ComfyUI-Gemini && pip install -r requirements.txt
6- cd ComfyUI/custom_nodes && git clone https://github.com/kijai/ComfyUI-Marigold --recursive && cd ComfyUI-Marigold && pip install -r requirements.txt
7- cd ComfyUI/custom_nodes && git clone https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92 --recursive
8- cd ComfyUI/custom_nodes && git clone https://github.com/Fannovel16/comfyui_controlnet_aux --recursive && cd comfyui_controlnet_aux && pip install -r requirements.txt
9- cd ComfyUI/models/controlnet && wget -O control-lora-canny-rank256.safetensors https://huggingface.co/stabilityai/control-lora/resolve/main/control-LoRAs-rank256/control-lora-canny-rank256.safetensors
10- cd ComfyUI/models/controlnet && wget -O control-lora-depth-rank256.safetensors https://huggingface.co/stabilityai/control-lora/resolve/main/control-LoRAs-rank256/control-lora-depth-rank256.safetensors
11- cd ComfyUI/models/checkpoints && wget -O dreamshaperXL_v21TurboDPMSDE.safetensors https://civitai.com/api/download/models/351306
12- cd ComfyUI/models/loras && wget -O StudioGhibli.Redmond-StdGBRRedmAF-StudioGhibli.safetensors https://huggingface.co/artificialguybr/StudioGhibli.Redmond-V2/resolve/main/StudioGhibli.Redmond-StdGBRRedmAF-StudioGhibli.safetensors
13environment_variables: {}
14external_package_dirs: []
15model_metadata: {}
16model_name: ComfyUI Anime Pet Style Transfer
17python_version: py310
18requirements:
19Â - websocket-client==1.6.4
20Â - accelerate==0.23.0
21Â - opencv-python
22resources:
23Â accelerator: A100
24Â use_gpu: true
25secrets: {}
26system_packages:
27Â - wget
28Â - ffmpeg
29Â - libgl1-mesa-glx
Step 2: Modifying the ComfyUI workflow to an API-compatible format
ComfyUI workflows can be run on Baseten by exporting them in an API format. Check out our blog on how to serve ComfyUI models behind an API endpoint if you need help converting your workflow accordingly.
For this tutorial, the workflow file can be copied from here. This workflow has two inputs: a prompt and an image. We can specify those variables inside our workflow JSON file using the handlebars template {{prompt}}
and {{input_image}}
.
That’s it! We can now deploy our ComfyUI workflow to Baseten!
Step 3: Deploying your ComfyUI workflow to Baseten
To deploy our workflow to Baseten, make sure you have the truss
Python package installed.Â
pip install truss --upgrade
With truss_examples/comfyui_truss
as the root directory, we can run the following command to deploy to Baseten:
truss push --publish
If prompted, paste in your Baseten API key. This command will package your Truss and deploy it onto Baseten’s cloud. The Docker container will be built and then deployed as an API endpoint.
Now we’ll run our first inference!
Running inference on your ComfyUI API endpoint
Once your model has deployed and is in the active
state, you can call the API endpoint like this:
1import requests
2import os
3import base64
4from PIL import Image
5from io import BytesIO
6
7# Replace the empty string with your model id below
8model_id = ""
9baseten_api_key = os.environ["BASETEN_API_KEY"]
10BASE64_PREAMBLE = "data:image/png;base64,"
11
12def pil_to_b64(pil_img):
13Â Â Â buffered = BytesIO()
14Â Â Â pil_img.save(buffered, format="PNG")
15Â Â Â img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
16Â Â Â return img_str
17
18def b64_to_pil(b64_str):
19Â Â Â return Image.open(BytesIO(base64.b64decode(b64_str.replace(BASE64_PREAMBLE, ""))))
20
21values = {
22Â "prompt": "american Shorthair",
23Â "input_image": {"type": "image", "data": pil_to_b64(Image.open("/path/to/cat.png"))}
24}
25
26resp = requests.post(
27Â Â Â f"https://model-{model_id}.api.baseten.co/production/predict",
28Â Â Â headers={"Authorization": f"Api-Key {baseten_api_key}"},
29Â Â Â json={"workflow_values": values}
30)
31
32res = resp.json()
33results = res.get("result")
34
35for item in results:
36Â Â Â if item.get("format") == "png":
37Â Â Â Â Â Â Â data = item.get("data")
38Â Â Â Â Â Â Â img = b64_to_pil(data)
39Â Â Â Â Â Â Â img.save(f"pet-style-transfer-1.png")
The API call returns result
, which is a list of images the model returns. The data
for each item in result
is the base64 string representation of the output image. We convert this base64 string to a PIL object and save it as an image.
Here are some sample inputs and their corresponding anime-stylized outputs:
That’s all it takes to run ComfyUI workflows as API endpoints on Baseten. Now, with the new build_commands
feature, it's even easier to add your own custom nodes and model checkpoints.Â
We’re dedicated to building an excellent developer experience around ComfyUI deployments. Try building your own custom ComfyUI workflow and run it as a production-grade API service, or try launching a sample workflow from our model library — either way, be sure to tell us what you think!