ComfyUI Workflow to API: A Comprehensive Guide in Python

ComfyUI is a powerful tool for creating image generation workflows. However, its true potential can be unlocked by converting these workflows into APIs, allowing for dynamic processing of user input and wider application integration. This guide will walk you through the process of transforming your ComfyUI workflow into a functional API.

Setting Up the API

The first step is to establish a connection with ComfyUI's WebSocket interface. This allows for real-time updates on the workflow's progress.

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

Key API Endpoints

Queueing a Prompt

To initiate the workflow, we need to send a prompt to ComfyUI:

def queue_prompt(prompt, client_id, server_address):
    p = {"prompt": prompt, "client_id": client_id}
    headers = {'Content-Type': 'application/json'}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request(f"http://{server_address}/prompt", data=data, headers=headers)
    return json.loads(urllib.request.urlopen(req).read())

Retrieving History

After processing, we can fetch the workflow's history:

def get_history(prompt_id, server_address):
    with urllib.request.urlopen(f"http://{server_address}/history/{prompt_id}") as response:
        return json.loads(response.read())

Fetching Generated Images

To retrieve the generated images:

def get_image(filename, subfolder, folder_type, server_address):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen(f"http://{server_address}/view?{url_values}") as response:
        return response.read()

Workflow Manipulation

To make your workflow dynamic, you'll need to modify certain aspects of it programmatically:

def prompt_to_image(workflow, positive_prompt, negative_prompt='', save_previews=False):
    prompt = json.loads(workflow)
    id_to_class_type = {id: details['class_type'] for id, details in prompt.items()}
    k_sampler = [key for key, value in id_to_class_type.items() if value == 'KSampler'][0]
    prompt[k_sampler]['inputs']['seed'] = random.randint(10**14, 10**15 - 1)
    positive_input_id = prompt[k_sampler]['inputs']['positive'][0]
    prompt[positive_input_id]['inputs']['text'] = positive_prompt
    if negative_prompt:
        negative_input_id = prompt[k_sampler]['inputs']['negative'][0]
        prompt[negative_input_id]['inputs']['text'] = negative_prompt
    generate_image_by_prompt(prompt, './output/', save_previews)

Tracking Progress

Utilize the WebSocket connection to monitor the workflow's progress:

def track_progress(prompt, ws, prompt_id):
    node_ids = list(prompt.keys())
    finished_nodes = []
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'progress':
                data = message['data']
                print(f"K-Sampler Progress: Step {data['value']} of {data['max']}")
            elif message['type'] in ['execution_cached', 'executing']:
                data = message['data']
                if data['node'] not in finished_nodes:
                    finished_nodes.append(data['node'])
                print(f"Progress: {len(finished_nodes)}/{len(node_ids)} Tasks completed")
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break
    return

Handling Generated Images

After the workflow completes, retrieve and save the generated images:

def get_images(prompt_id, server_address, allow_preview=False):
    output_images = []
    history = get_history(prompt_id, server_address)[prompt_id]
    for node_id in history['outputs']:
        node_output = history['outputs'][node_id]
        if 'images' in node_output:
            for image in node_output['images']:
                if (allow_preview and image['type'] == 'temp') or image['type'] == 'output':
                    image_data = get_image(image['filename'], image['subfolder'], image['type'], server_address)
                    output_images.append({
                        'image_data': image_data,
                        'file_name': image['filename'],
                        'type': image['type']
                    })
    return output_images
 
def save_image(images, output_path, save_previews):
    for item in images:
        directory = os.path.join(output_path, 'temp/') if item['type'] == 'temp' and save_previews else output_path
        os.makedirs(directory, exist_ok=True)
        try:
            image = Image.open(io.BytesIO(item['image_data']))
            image.save(os.path.join(directory, item['file_name']))
        except Exception as e:
            print(f"Failed to save image {item['file_name']}: {e}")

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 these steps, you can transform your ComfyUI workflow into a versatile API. This opens up possibilities for integrating your image generation pipeline into various applications, allowing for dynamic input processing and real-time feedback. Whether you're building a web application, a mobile app, or a complex AI system, this API approach provides the flexibility and power to leverage ComfyUI's capabilities in diverse scenarios