快速入门
本页将带你一步步完成自定义节点的创建过程。
我们的示例将接收一批图片,并返回其中一张图片。最初,这个节点会返回平均颜色最亮的图片;随后我们会扩展它,支持多种选择标准,最后还会添加一些客户端代码。
本页假设你对 Python 或 Javascript 的了解很少。
编写基础节点
前置条件
环境搭建
回答几个问题后,你会得到一个新的目录。
定义节点
将以下代码添加到 src/nodes.py
末尾:
自定义节点通过 Python 类定义,必须包含以下四项:CATEGORY
(指定新节点在添加节点菜单中的位置)、INPUT_TYPES
(类方法,定义节点输入,详见后文)、RETURN_TYPES
(定义节点输出)、FUNCTION
(节点执行时调用的函数名)。
IMAGE
(单数),即使我们期望接收一批图片并只返回一张。在 Comfy 中,IMAGE
代表图片批量,单张图片视为批量大小为 1。主函数
主函数 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 才能看到更改。
添加选项
这个节点可能有点无聊,所以我们可以加一些选项;比如一个小部件,让你选择最亮、最红、最绿或最蓝的图片。将你的 INPUT_TYPES
修改为:
然后更新主函数。我们用一个很简单的“最红”定义,即像素的平均 R 值除以三色平均值。所以:
调整 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 服务器,重新启动,刷新网页,运行你的工作流。