告别等待：ComfyUI Websocket实时通信技术详解

2026-02-05 04:59:45作者：邓越浪Henry

你是否还在为AI图像生成过程中的漫长等待而烦恼？是否希望能实时看到图像生成的每一步进展？ComfyUI的Websocket（套接字）技术彻底改变了这一现状，让原本需要等待完整生成才能查看结果的流程转变为实时流式体验。本文将深入解析ComfyUI的Websocket实现原理，带你掌握如何利用这一技术构建流畅的实时交互体验。

实时通信架构概览

ComfyUI采用了基于WebSocket协议的双向通信架构，实现了客户端与服务器之间的实时数据传输。这一架构主要由以下核心组件构成：

WebSocket服务器：位于server.py中，负责建立和管理WebSocket连接
图像传输节点：如custom_nodes/websocket_image_save.py，处理图像数据的实时发送
前端接收处理：负责在浏览器中实时显示接收到的图像数据

通信流程示意图

sequenceDiagram
    participant 客户端
    participant WebSocket服务器
    participant 图像生成节点
    participant 进度条系统
    
    客户端->>WebSocket服务器: 建立连接(握手)
    WebSocket服务器-->>客户端: 连接确认
    客户端->>WebSocket服务器: 发送生成任务
    WebSocket服务器->>图像生成节点: 执行生成任务
    
    loop 生成过程中
        图像生成节点->>进度条系统: 更新进度
        进度条系统->>WebSocket服务器: 进度数据
        WebSocket服务器->>客户端: 发送进度更新
        
        图像生成节点-->>WebSocket服务器: 中间图像结果
        WebSocket服务器-->>客户端: 流式传输图像数据
    end
    
    图像生成节点->>WebSocket服务器: 生成完成
    WebSocket服务器->>客户端: 发送最终结果和状态
    客户端->>WebSocket服务器: 断开连接

WebSocket服务器实现解析

WebSocket服务器的核心实现位于server.py文件中，主要通过websocket_handler函数处理客户端连接和消息交互。

连接建立与管理

服务器在接收到WebSocket连接请求时，会执行以下步骤：

为客户端生成唯一会话ID（SID）
存储WebSocket连接对象和客户端元数据
发送初始状态信息，包括队列状态和会话ID

关键代码实现如下：

@routes.get('/ws')
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    sid = request.rel_url.query.get('clientId', '')
    if sid:
        # 重用现有会话，移除旧连接
        self.sockets.pop(sid, None)
    else:
        sid = uuid.uuid4().hex

    # 存储WebSocket连接
    self.sockets[sid] = ws
    # 存储元数据
    self.sockets_metadata[sid] = {"feature_flags": {}}

    try:
        # 向新客户端发送初始状态
        await self.send("status", {"status": self.get_queue_info(), "sid": sid}, sid)
        # ...处理消息循环...

消息处理机制

服务器实现了异步消息处理循环，能够同时处理多个客户端连接：

async for msg in ws:
    if msg.type == aiohttp.WSMsgType.ERROR:
        logging.warning('ws connection closed with exception %s' % ws.exception())
    elif msg.type == aiohttp.WSMsgType.TEXT:
        try:
            data = json.loads(msg.data)
            # 处理客户端发送的消息
            # ...
        except json.JSONDecodeError:
            logging.warning(f"Invalid JSON received from client {sid}: {msg.data}")

实时图像传输节点

custom_nodes/websocket_image_save.py提供了专门的节点，用于将生成过程中的图像通过WebSocket实时发送到客户端。

节点实现原理

该节点通过重写图像保存逻辑，将图像数据直接通过WebSocket发送，而非仅保存到磁盘：

class SaveImageWebsocket:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"images": ("IMAGE", ),}}

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "api/image"

    def save_images(self, images):
        pbar = comfy.utils.ProgressBar(images.shape[0])
        step = 0
        for image in images:
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            # 通过进度条系统发送图像数据
            pbar.update_absolute(step, images.shape[0], ("PNG", img, None))
            step += 1

        return {}