How to Use the ComfyUI WebSocket API for Real-Time Image Generation

ComfyUI exposes a bidirectional WebSocket endpoint at /ws that streams execution events and binary image data, enabling clients to receive generated images in real-time without polling the filesystem.

ComfyUI, the open-source node-based Stable Diffusion interface, provides a powerful WebSocket API for real-time communication between the server and client. By leveraging the /ws endpoint, developers can build live-preview applications, external integrations, and automated workflows that receive images instantly as they are generated. This guide explains the protocol implementation as found in the Comfy-Org/ComfyUI repository and provides practical Python examples for consuming the stream.

Understanding the WebSocket Protocol

The server-side implementation lives in server.py, where a PromptServer instance creates a WebSocketResponse and stores each connection in self.sockets. When a client connects to ws://HOST:PORT/ws?clientId=UUID, the server initiates a handshake and begins streaming events.

Connection Handshake and Feature Flags

Upon connection, the client receives an initial status message (type: "status") containing current queue statistics. The server then negotiates feature flags:

1. The client may send a "feature_flags" payload as its first JSON message
2. The server stores these flags under self.sockets_metadata[sid]["feature_flags"]
3. The server replies with its own capabilities via feature_flags.get_server_features()

This handshake allows the server to enable optional features like progress reporting based on client capabilities.

Binary Event Types and Message Structure

Binary data transmission uses a specific framing protocol defined in protocol.py. The BinaryEventTypes enum defines event codes:

Event	Code	Description
PREVIEW_IMAGE	1	JPEG/PNG preview image during generation
UNENCODED_PREVIEW_IMAGE	2	Raw image data without extra framing
PREVIEW_IMAGE_WITH_METADATA	4	Image data with length-prefixed JSON metadata

Binary messages follow an 8-byte header structure: 4 bytes for the event type (little-endian integer) followed by 4 reserved bytes. The actual payload begins at byte 8.

JSON Message Types

During execution, the server sends structured JSON messages:

Message Type	Payload	Purpose
executing	{ "prompt_id": "...", "node": "<node_id>" }	Signals node execution start; node is null when the workflow completes
status	{ "status": … }	Queue statistics (running/queued counts)
feature_flags	{ … }	Feature negotiation results
progress	{ "value": … }	Optional percent-complete updates

Building a Real-Time Client

To capture generated images without writing to disk, include the SaveImageWebsocket node in your workflow. This node emits raw image bytes through the WebSocket connection rather than saving to the output folder.

Receiving Images with SaveImageWebsocket

The following Python client demonstrates complete workflow execution with real-time image capture. It handles both JSON status messages and binary image payloads:

import websocket               # pip install websocket-client

import uuid
import json
import urllib.request
import urllib.parse
from PIL import Image
import io

SERVER = "127.0.0.1:8188"
CLIENT_ID = str(uuid.uuid4())

def queue_prompt(prompt):
    """Submit a prompt and return the server‑generated prompt_id."""
    payload = {"prompt": prompt, "client_id": CLIENT_ID}
    data = json.dumps(payload).encode()
    req = urllib.request.Request(f"http://{SERVER}/prompt", data=data)
    return json.loads(urllib.request.urlopen(req).read())

def get_images(ws, prompt):
    """Run a workflow and collect images sent by SaveImageWebsocket."""
    prompt_id = queue_prompt(prompt)["prompt_id"]
    images_by_node = {}
    current_node = None

    while True:
        msg = ws.recv()
        # ----- JSON messages -------------------------------------------------

        if isinstance(msg, str):
            data = json.loads(msg)
            if data["type"] == "executing":
                info = data["data"]
                if info["prompt_id"] != prompt_id:
                    continue      # ignore other clients

                if info["node"] is None:      # workflow finished

                    break
                current_node = info["node"]
            elif data["type"] == "feature_flags":
                # optional – ignore or store server features

                pass
        # ----- Binary image payload -----------------------------------------

        else:
            # Binary events are: 4‑byte header + image bytes

            # For SaveImageWebsocket the header is 0x00000000 (unused)

            image_bytes = msg[8:]               # skip 4‑byte event + 4‑byte reserved

            if current_node == "save_image_websocket_node":
                images_by_node.setdefault(current_node, []).append(image_bytes)

    # Convert raw bytes to Pillow images (optional)

    for node, blobs in images_by_node.items():
        images_by_node[node] = [Image.open(io.BytesIO(b)) for b in blobs]
    return images_by_node

# -------------------------------------------------------------------------

# Example workflow (JSON) – insert your own node IDs / parameters

# -------------------------------------------------------------------------

workflow = {
    "3": {"class_type": "KSampler", "inputs": {"cfg": 8, "denoise": 1,
            "latent_image": ["5", 0], "model": ["4", 0],
            "negative": ["7", 0], "positive": ["6", 0],
            "sampler_name": "euler", "scheduler": "normal",
            "seed": 5, "steps": 20}},
    "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "v1-5-pruned-emaonly.safetensors"}},
    "5": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1, "height": 512, "width": 512}},
    "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["4", 1], "text": "masterpiece best quality man"}},
    "7": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["4", 1], "text": "bad hands"}},
    "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
    "save_image_websocket_node": {"class_type": "SaveImageWebsocket",
                                 "inputs": {"images": ["8", 0]}}
}

# -------------------------------------------------------------------------

# Run the client

# -------------------------------------------------------------------------

ws = websocket.WebSocket()
ws.connect(f"ws://{SERVER}/ws?clientId={CLIENT_ID}")
result_images = get_images(ws, workflow)
ws.close()

# Show the first generated image (optional)

if result_images:
    first_image = next(iter(result_images.values()))[0]
    first_image.show()

Key implementation details:

• WebSocket URL – Must include the clientId query parameter for the server to associate messages with your connection
• Binary handling – Slice msg[8:] to remove the 8-byte header (4 bytes event type + 4 bytes reserved)
• Node tracking – The executing message indicates which node is currently running, allowing you to associate binary image data with the correct workflow step
• SaveImageWebsocket – Replace any standard SaveImage node with this variant to receive raw bytes instead of filesystem writes

Monitoring Execution Progress

For use cases requiring only status updates without image data, implement a minimal monitor that tracks node execution:

import websocket, json, uuid, urllib.request

SERVER = "127.0.0.1:8188"
CID = str(uuid.uuid4())

def queue(prompt):
    payload = {"prompt": prompt, "client_id": CID}
    req = urllib.request.Request(f"http://{SERVER}/prompt",
                                 data=json.dumps(payload).encode())
    return json.loads(urllib.request.urlopen(req).read())["prompt_id"]

def monitor(prompt):
    ws = websocket.WebSocket()
    ws.connect(f"ws://{SERVER}/ws?clientId={CID}")

    prompt_id = queue(prompt)
    while True:
        msg = ws.recv()
        if isinstance(msg, str):
            data = json.loads(msg)
            if data["type"] == "executing":
                d = data["data"]
                if d["prompt_id"] == prompt_id and d["node"] is None:
                    print("Workflow finished")
                    break
                print(f"Running node: {d['node']}")
    ws.close()

This pattern is useful for progress bars, external orchestration systems, or logging node execution times.

Server-Side Architecture

Understanding the server implementation helps debug connection issues and optimize client behavior:

• server.py – Contains the PromptServer class that manages the /ws route. It maintains the self.sockets dictionary mapping client IDs to WebSocket objects and handles the send_image and send_image_with_metadata methods for binary transmission.
• protocol.py – Defines the BinaryEventTypes enum used to tag binary messages. The server imports these constants when encoding preview images or websocket-saved images.
• script_examples/websockets_api_example_ws_images.py – Official reference implementation demonstrating SaveImageWebsocket usage with complete error handling.
• script_examples/websockets_api_example.py – Simpler example showing the traditional polling approach via the /history endpoint, useful for comparing WebSocket versus HTTP polling architectures.

Summary

• WebSocket endpoint – Connect to /ws?clientId=UUID to establish a bidirectional stream with the ComfyUI server.
• Binary protocol – Image data arrives with an 8-byte header (event type + reserved); slice at byte 8 to access raw JPEG/PNG bytes.
• SaveImageWebsocket node – Essential for receiving final images through the socket rather than reading from disk.
• Execution tracking – Monitor the executing JSON message type to determine when nodes start and when the workflow completes (node: null).
• Feature negotiation – The server supports capability negotiation via feature_flags messages for enabling optional progress reporting.

Frequently Asked Questions

How do I receive live previews during generation?

The server automatically sends PREVIEW_IMAGE (type 1) binary events during sampling. These arrive on the same WebSocket connection with the standard 8-byte header. Decode them by stripping the header and loading the remaining bytes as a JPEG or PNG image. Previews are generated by the PreviewImage node or automatically when using certain samplers that support latent preview encoding.

What is the difference between SaveImageWebsocket and regular SaveImage?

SaveImageWebsocket emits raw image bytes directly through the WebSocket connection using the binary protocol, bypassing the filesystem entirely. The standard SaveImage node writes images to the output folder on disk and returns a URL. Use SaveImageWebsocket for real-time integrations where disk I/O is unnecessary or when building headless automation pipelines.

Why does my client receive binary messages with only 4 bytes?

The server may send control messages or empty previews. Always check the message length before slicing. Valid image payloads from SaveImageWebsocket include the 4-byte event code (typically 0 or 1) followed by 4 reserved bytes, then the actual image data. If msg[8:] is empty, the message represents a control signal or empty preview frame that should be ignored.

Can multiple clients connect to the same workflow execution?

Yes, but each client must use a unique clientId UUID. The server broadcasts executing messages to all connected sockets, so clients must filter messages by prompt_id to identify their own workflow results. The queue_prompt function returns a unique prompt_id that clients should match against incoming executing events to ignore other users' jobs.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}