Software 2.0

ComfyUI自定义节点开发演练

admin -

本指南将逐步指导你创建自定义节点的过程，该节点获取一批图像并返回其中一张图像。最初，节点将返回平均颜色最浅的图像；然后我们将对其进行扩展，使其具有一系列选择标准，最后添加一些客户端代码。

本页假设你对 Python 或 Javascript 知之甚少。

在本指南之后，深入了解服务器端代码、客户端代码或客户端-服务器通信的详细信息。

1、基本节点

此自定义节点的所有代码将位于单个目录中。因此，首先在 ComfyUI 文件夹中找到 custom_nodes 目录，然后在其中创建一个新目录，例如命名为 image_selector。这个新目录是与新自定义节点相关的所有代码的基本目录。

1.1 Python代码框架

自定义节点的基本结构将在后面详细描述。我们从最基本的必需品开始：

class ImageSelector:
    CATEGORY = "example"
    @classmethod    
    def INPUT_TYPES(s):
        return { "required":  { "images": ("IMAGE",), } }
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "choose_image"

自定义节点是一个 Python 类，必须包含以下四个内容：

  • CATEGORY，指定自定义节点在添加新节点菜单中的位置，
  • INPUT_TYPES，定义节点将采用哪些输入的类方法（稍后查看返回的字典的详细信息），
  • RETURN_TYPES，定义节点将产生哪些输出，
  • FUNCTION，执行节点时将调用的函数的名称。

请注意，输入和输出的数据类型是 IMAGE（单数），即使我们期望接收一批图像并只返回一个。在 Comfy 中，IMAGE 表示图像批次，单个图像被视为大小为 1 的批次。

1.2 添加主函数

主函数 choose_image 接收 INPUT_TYPES 中定义的命名参数，并返回 RETURN_TYPES 中定义的元组。由于我们处理的图像在内部存储为 torch.Tensor，

import torch

然后将函数添加到你的类中。图像的数据类型为 torch.Tensor，形状为 [B,H,W,C]，其中 B 是批次大小，C 是通道数 - RGB 为 3。如果我们迭代这样的张量，我们将得到一系列形状为 [H,W,C] 的 B 张量。 .flatten() 方法将其转换为长度为 HWC 的一维张量， torch.mean() 取平均值， . item() 将单值张量转换为 Python 浮点数。

    def choose_image(self, images):
        brightness = list(torch.mean(image.flatten()).item() for image in images)
        brightest = brightness.index(max(brightness))
        result = images[brightest].unsqueeze(0)
        return (result,)

最后两行的注释：

  • images[brightest] 将返回形状为 [H,W,C] 的张量。 unsqueeze 用于插入一个（长度为 1）维度，在本例中为维度 0，从而得到  [B,H,W,C] 其中 B=1：单个图像。
  • 在 return (result,) 中，结尾的逗号对于确保返回元组至关重要。​

1.3 部署节点

为了使 Comfy 识别新节点，我们需要将目录 image_selector 转换为 Python 模块，方法是添加 __init__.py，如下所示：

from .image_selector_node import ImageSelector

NODE_CLASS_MAPPINGS = {
    "Image Selector" : ImageSelector,
}

__all__ = ['NODE_CLASS_MAPPINGS']

这里我们只是导出 NODE_CLASS_MAPPINGS，它为每个新的自定义节点提供了一个映射到类的唯一名称。

​1.4 运行 Comfy

启动（或重新启动）Comfy 服务器，你应该在自定义节点列表中看到如下一行：

0.0 seconds: [your path]\ComfyUI\custom_nodes\image_selector

在浏览器中重新加载 Comfy 页面，在“添加节点”菜单的示例下，你将找到 image_selector。如果没有，请在 Python 控制台输出中查找错误！

2、添加一些选项

该节点可能有点无聊，所以我们可能会添加一些选项；一个小部件，允许你选择最亮的图像，或者最红、最蓝或最绿的图像。编辑 Python带啊 以添加另一个输入，因此 INPUT_TYPES 看起来像：

    @classmethod    
    def INPUT_TYPES(s):
        return { "required":  { "images": ("IMAGE",), 
                                "mode": (["brightest", "reddest", "greenest", "bluest"],)} }

然后更新主函数。我们将使用一个相当简单的“reddest”定义，即像素的平均 R 值除以所有三种颜色的平均值。所以：

    def choose_image(self, images, mode):
        batch_size = images.shape[0]
        brightness = list(torch.mean(image.flatten()).item() for image in images)
        if (mode=="brightest"):
            scores = brightness
        else:
            channel = 0 if mode=="reddest" else (1 if mode=="greenest" else 2)
            absolute = list(torch.mean(image[:,:,channel].flatten()).item() for image in images)
            scores = list( absolute[i]/(brightness[i]+1e-8) for i in range(batch_size) )
        best = scores.index(max(scores))
        result = images[best].unsqueeze(0)
        return (result,)

3、调整 UI

也许我们想要一些视觉反馈，所以让我们发送一条小短信来显示。

​3.1 从服务器发送消息

这需要在 Python 代码中添加两行：

from server import PromptServer

并且在 choose_image 方法的末尾添加一行以向前端发送消息。 send_sync 输入参数为消息类型（应是唯一的）和字典：

        PromptServer.instance.send_sync("example.imageselector.textmessage", {"message":f"Picked image {best+1}"})
        return (result,)

3.2 编写客户端扩展

要向客户端添加一些 Javascript，请在​​自定义节点目录中创建一个子目录 js，并修改 __init__.py 的末尾以通过导出 WEB_DIRECTORY 将其告知 Comfy：

WEB_DIRECTORY = "./js"
__all__ = ['NODE_CLASS_MAPPINGS', 'WEB_DIRECTORY']

客户端扩展在 js 中保存为 .js 文件子目录，因此使用以下代码创建 image_selector/js/image_selector.js。有关详细信息，请参阅客户端编码。

import { app } from "../../../scripts/app.js";
import { api } from "../../../scripts/api.js";

app.registerExtension({
	name: "example.imageselector",
    async setup() {
        function messageHandler(event) { alert(event.detail.message); }
        api.addEventListener("example.imageselector.textmessage", messageHandler);
    },
})

我们所做的就是注册一个扩展，并在其 setup() 方法中，为我们发送的消息类型添加一个侦听器，并读取我们发送的存储在 event.detail 中字典。

停止 Comfy 服务器，重新启动它，重新加载网页，然后运行你的工作流程。

原文链接：Comfy Custom Nodes Walkthrough

汇智网翻译整理，转载请标明出处