编写基础节点
前置条件
环境搭建
定义节点
将以下代码添加到src/nodes.py 末尾:
src/nodes.py
自定义节点的基本结构详见这里。
CATEGORY(指定新节点在添加节点菜单中的位置)、INPUT_TYPES(类方法,定义节点输入,详见后文)、RETURN_TYPES(定义节点输出)、FUNCTION(节点执行时调用的函数名)。
主函数
主函数choose_image 会收到在 INPUT_TYPES 中定义的命名参数,并返回一个与 RETURN_TYPES 匹配的 tuple。由于我们处理的是图片,图片在内部以 torch.Tensor 存储,
[B,H,W,C] 的 torch.Tensor,其中 B 是批量大小,C 是通道数(RGB 为 3)。遍历该张量会得到 B 个形状为 [H,W,C] 的张量。.flatten() 方法将其变为一维张量,长度为 H*W*C,torch.mean() 求均值,.item() 将单值张量转为 Python 浮点数。
images[brightest]返回形状为[H,W,C]的张量。unsqueeze用于在第 0 维插入一个长度为 1 的维度,得到[B,H,W,C],其中B=1,即单张图片。return (result,)末尾的逗号很重要,确保返回的是元组。
注册节点
要让 Comfy 识别新节点,必须在包级别可用。修改src/nodes.py 末尾的 NODE_CLASS_MAPPINGS 变量。你需要重启 ComfyUI 才能看到更改。
src/nodes.py
关于 ComfyUI 如何发现和加载自定义节点的详细说明,请参见节点生命周期文档。
添加选项
这个节点可能有点无聊,所以我们可以加一些选项;比如一个小部件,让你选择最亮、最红、最绿或最蓝的图片。将你的INPUT_TYPES 修改为:
调整 UI
也许我们想要一些可视化反馈,所以让我们发送一条文本消息进行显示。从服务器发送消息
只需在 Python 代码中添加两行:choose_image 方法末尾添加一行,将消息发送到前端(send_sync 需要一个唯一的消息类型和一个字典):
编写客户端扩展
要为客户端添加 Javascript,在你的自定义节点目录下创建web/js 子目录,并在 __init__.py 末尾导出 WEB_DIRECTORY:
.js 文件保存在 web/js 子目录下,所以创建 image_selector/web/js/imageSelector.js,内容如下。(更多内容见 客户端开发)
setup() 方法中为我们发送的消息类型添加监听器。它会读取我们发送的字典(存储在 event.detail 中)。
停止 Comfy 服务器,重新启动,刷新网页,运行你的工作流。
完整示例
完整示例见这里。你可以下载示例工作流 JSON 文件 或在下方查看: