如何实现ComfyUI自定义节点的数据传递与处理？

Question

在开发ComfyUI自定义节点时，如何正确实现节点间的数据传递与处理是一个核心问题。常见的疑问包括：如何定义输入输出端口的数据类型？如何在节点函数中接收和处理上游节点传递的数据？以及如何将处理结果传递给下游节点？此外，对于复杂数据结构如图像、模型或张量，如何确保数据在流程中正确序列化与反序列化？理解ComfyUI的节点基类方法（如INPUT_TYPES、RETURN_TYPES、FUNCTION）及其执行机制是关键。如何利用节点的__init__方法初始化状态数据？如何处理异步或批量数据？这些问题直接影响节点的稳定性与扩展性。

Answer 1

一、ComfyUI节点数据传递机制概述

在ComfyUI中，节点是构建可视化流程图的核心单元。每个节点通过定义输入输出端口、执行函数逻辑以及状态管理，实现与其他节点的数据交互。理解其底层机制是开发稳定、高效节点的前提。

• 节点数据传递基于Python字典结构
• 数据在节点间以命名端口为单位进行传递
• 执行流程由ComfyUI调度器控制，支持同步与异步处理

二、输入输出端口定义

使用类变量 INPUT_TYPES 和 RETURN_TYPES 定义节点的输入输出类型，是节点间通信的基础。

class MyCustomNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": "IMAGE",
                "scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0})
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process_image"

常见数据类型包括：

类型名	描述	示例值
IMAGE	图像张量（PyTorch Tensor）	torch.Tensor([3, 512, 512])
MODEL	模型对象（如Diffusion模型）	diffusers.StableDiffusionPipeline
STRING	字符串类型	"prompt text"
INT	整数类型	42

三、节点函数的执行机制

通过 FUNCTION 类变量指定执行函数，该函数将接收上游节点传递的参数，并返回结果给下游节点。

def process_image(self, image, scale):
    # 处理图像逻辑
    scaled_image = F.interpolate(image.unsqueeze(0), scale_factor=scale).squeeze(0)
    return (scaled_image,)

执行流程如下：

graph TD A[Node A Output] --> B[Node B Input] B --> C[Node B FUNCTION] C --> D[Node B Output] D --> E[Node C Input]

• ComfyUI自动将上游输出与下游输入按名称匹配
• 函数返回值顺序需与RETURN_TYPES定义一致
• 支持返回多个结果，如：RETURN_TYPES = ("IMAGE", "MASK")

四、复杂数据结构的序列化与反序列化

ComfyUI支持图像、模型等复杂数据类型，但需注意其在流程中如何传递：

• 图像数据通常为PyTorch Tensor，直接传递即可
• 模型对象应确保为可序列化结构（如仅传递模型配置或引用）
• 对于不可序列化的对象，建议使用 __getstate__ 和 __setstate__ 方法自定义序列化逻辑

class MyModelNode:
    def __init__(self):
        self.model = load_model("my_model.pth")

    def __getstate__(self):
        return {}  # 避免序列化整个模型

    def __setstate__(self, state):
        self.model = load_model("my_model.pth")  # 重新加载

五、节点状态初始化与生命周期管理

使用 __init__ 方法初始化节点状态数据，适用于缓存、配置加载等场景。

class CacheNode:
    def __init__(self):
        self.cache = {}

    def store_data(self, key, value):
        self.cache[key] = value

• 状态数据在节点执行周期内保持有效
• 适合用于缓存中间结果、避免重复计算
• 注意线程安全和异步执行时的状态一致性

六、异步与批量数据处理

ComfyUI支持异步执行和批量处理，需在节点中进行相应配置：

1. 在节点类中定义 FUNCTION = "async_process"
2. 使用 async def 定义异步函数
3. 批量处理可通过循环或向量化操作实现

async def async_process(self, images, batch_size=4):
    results = []
    for i in range(0, len(images), batch_size):
        batch = images[i:i+batch_size]
        processed = await process_batch_async(batch)
        results.extend(processed)
    return (results,)

• 异步节点需确保线程安全
• 批量处理可显著提升性能，但需考虑内存限制