ComfyUI Workflow into API with an Actual Example in Python

ComfyUI is an advanced and powerful graphical user interface for creating complex image generation workflows, particularly in the field of AI and machine learning. While it's highly effective for visual experimentation, integrating these workflows into applications requires an API-based approach. This tutorial provides a step-by-step guide to converting your ComfyUI workflow into a functional API using Python, complete with actual code examples.

Prerequisites
Understanding the Workflow
Setting Up the Environment
Creating the API Endpoints
Establishing WebSocket Connections
Modifying the Workflow Dynamically
Tracking Progress
Retrieving and Serving Images
Putting It All Together
Conclusion

Prerequisites

Before we begin, ensure you have the following:

Python 3.6 or higher installed.
ComfyUI installed and running on your machine.
Basic understanding of Python programming.
Familiarity with RESTful APIs and WebSockets.

Understanding the Workflow

A ComfyUI workflow consists of interconnected nodes that define the image generation process. Each node performs a specific function, such as applying a filter or generating noise. By converting this workflow into an API, you enable external applications to interact with it programmatically.

Setting Up the Environment

First, create a new Python project or virtual environment to manage dependencies. Install the required packages:

pip install websocket-client flask flask-socketio

websocket-client: For handling WebSocket communication.
Flask: A lightweight web framework to create API endpoints.
Flask-SocketIO: To integrate WebSocket support in Flask.

Creating the API Endpoints

We'll use Flask to create RESTful API endpoints that allow clients to interact with the ComfyUI workflow.

Initialize Flask App

Create a new file app.py and set up the basic Flask application:

from flask import Flask, request, jsonify
from flask_socketio import SocketIO
 
app = Flask(__name__)
socketio = SocketIO(app)

Endpoint to Queue a Prompt

Create an endpoint to queue a new prompt:

@app.route('/api/queue_prompt', methods=['POST'])
def queue_prompt():
    data = request.get_json()
    prompt = data.get('prompt')
    client_id = data.get('client_id')
 
    # Logic to send the prompt to ComfyUI
    # ...
 
    return jsonify({'status': 'Prompt queued successfully', 'prompt_id': 'some_prompt_id'}), 200

Endpoint to Retrieve History

Create an endpoint to retrieve the workflow's history:

@app.route('/api/history/<prompt_id>', methods=['GET'])
def get_history(prompt_id):
    # Logic to fetch history from ComfyUI
    # ...
 
    return jsonify({'history': '...'}), 200

Establishing WebSocket Connections

To receive real-time updates, we'll establish a WebSocket connection with ComfyUI.

Open WebSocket Connection

Create a utility function websocket_utils.py:

import websocket
import uuid
 
def open_websocket_connection():
    server_address = '127.0.0.1:8188'
    client_id = str(uuid.uuid4())
    ws = websocket.WebSocket()
    ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
    return ws, server_address, client_id

Modifying the Workflow Dynamically

To make the workflow respond to dynamic inputs, we'll modify its parameters programmatically.

import json
import random
 
def modify_workflow(workflow_json, positive_prompt, negative_prompt=''):
    workflow = json.loads(workflow_json)
    
    # Find the KSampler node
    k_sampler_node = next(
        (node_id for node_id, node in workflow.items() if node['class_type'] == 'KSampler'), 
        None
    )
 
    if k_sampler_node:
        # Update the seed with a random number
        workflow[k_sampler_node]['inputs']['seed'] = random.randint(1e14, 1e15 - 1)
        
        # Update the positive prompt
        positive_node_id = workflow[k_sampler_node]['inputs']['positive'][0]
        workflow[positive_node_id]['inputs']['text'] = positive_prompt
        
        # Update the negative prompt if provided
        if negative_prompt:
            negative_node_id = workflow[k_sampler_node]['inputs']['negative'][0]
            workflow[negative_node_id]['inputs']['text'] = negative_prompt
 
    return json.dumps(workflow)

Tracking Progress

Use the WebSocket connection to track the progress of the image generation.

import threading
 
def track_progress(ws, prompt_id):
    def on_message():
        while True:
            try:
                message = ws.recv()
                data = json.loads(message)
                if data.get('type') == 'progress' and data.get('data', {}).get('prompt_id') == prompt_id:
                    progress = data['data']
                    print(f"Progress: {progress['value']} / {progress['max']}")
                if data.get('type') == 'execution_complete' and data.get('data', {}).get('prompt_id') == prompt_id:
                    print("Image generation completed.")
                    break
            except websocket.WebSocketConnectionClosedException:
                break
 
    thread = threading.Thread(target=on_message)
    thread.start()

Retrieving and Serving Images

Once the image generation is complete, retrieve and serve the images through your API.

Fetch Generated Images

import requests
 
def get_generated_images(server_address, prompt_id):
    response = requests.get(f"http://{server_address}/history/{prompt_id}")
    history = response.json()
 
    images = []
    for node_output in history.get('outputs', {}).values():
        images.extend(node_output.get('images', []))
 
    return images

Serve Images via API

Add an endpoint to serve the images:

@app.route('/api/images/<prompt_id>', methods=['GET'])
def serve_images(prompt_id):
    images = get_generated_images('127.0.0.1:8188', prompt_id)
    # Logic to serve images, e.g., returning URLs or base64-encoded images
    return jsonify({'images': images}), 200

Putting It All Together

Here's how the entire flow works:

Client sends a POST request to /api/queue_prompt with the positive and negative prompts.
Server modifies the workflow using modify_workflow() and queues it in ComfyUI.
Server starts tracking progress using track_progress().
Once complete, Client can retrieve images by sending a GET request to /api/images/<prompt_id>.

Example Usage

Client Side

import requests
 
# Queue a new prompt
response = requests.post('http://localhost:5000/api/queue_prompt', json={
    'prompt': 'A beautiful sunset over the mountains',
    'client_id': 'your_client_id'
})
data = response.json()
prompt_id = data['prompt_id']
 
# Check progress (you can set up server-sent events or polling)
 
# Retrieve images
response = requests.get(f'http://localhost:5000/api/images/{prompt_id}')
images = response.json().get('images', [])

Or, You Can Use ComfyUIAsAPI.com's Service

If you don't want to deal with the hassle of setting up the API, you can use ComfyUIAsAPI.com's service.

All you need is to upload your ComfyUI workflow .json file and get a ready-to-use API. We also support SDKs for all the popular languages. You just need a few lines of code to integrate it into your project.

Conclusion

By following this step-by-step tutorial, you've transformed your ComfyUI workflow into a functional API using Python. This enables seamless integration of your image generation pipeline with other applications, allowing for dynamic input and real-time feedback. The flexibility of this approach empowers you to build more interactive and responsive AI-powered applications.