手把手教程

手把手教程

此页面将逐步引导您创建一个自定义节点，该节点接收一批图像并返回其中一幅图像。最初，节点将返回平均颜色最浅的图像；然后，我们将扩展它以具有多种选择标准，最后添加一些客户端代码。

假设您对 Python 或 Javascript 知之甚少。

基本节点

设置

此自定义节点的所有代码都将在一个目录中。因此，首先找到 ComfyUI 文件夹中的 custom_nodes 目录，并在其中创建一个新目录，命名为（例如）image_selector。此新目录是与新自定义节点相关的所有代码的基础目录。

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 的批处理。

添加主函数

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

import torch

然后将函数添加到您的类中。图像的数据类型是 torch.Tensor，其形状为 [B,H,W,C]，其中 B 是批处理大小，C 是通道数——对于 RGB 来说是 3。如果我们遍历这样的张量，将得到一系列形状为 [H,W,C] 的 B 张量。.flatten() 方法将其转换为一维张量，长度为 H*W*C，torch.mean() 计算平均值，.item() 将单个值张量转换为 Python float。

    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）维度，从而得到 [B,H,W,C]，其中 B=1：一幅单独的图像。
在 return (result,) 中，尾随逗号对于确保返回一个元组是必不可少的。

部署节点

为了让 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，为每个新自定义节点提供一个唯一名称，并映射到类。

运行 Comfy

启动（或重启）Comfy 服务器，您应该会在自定义节点列表中看到类似以下内容的行：

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

在浏览器中重新加载 Comfy 页面，在 Add Node 菜单下的 example 中，您会找到 image_selector。如果没有，请在 Python 控制台输出中查找错误！

添加一些选项

这个节点可能有点无聊，所以我们可以添加一些选项；一个小部件，允许您选择最亮的图像、最红的、最蓝的或最绿色的。编辑您的 Python 代码，添加另一个输入，使 INPUT_TYPES 看起来像：

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

然后更新主函数。我们将使用一个相当简单的“最红”的定义，即像素的平均 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,)

调整 UI

也许我们想要一点视觉反馈，因此让我们发送一条小文本消息以进行显示。

从服务器发送消息

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

from server import PromptServer

并在 choose_image 方法的末尾添加一行，以将消息发送到前端（send_sync 需要一个唯一的消息类型和一个字典）。

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

编写客户端扩展

要向客户端添加一些 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 服务器，再次启动它，重新加载网页，然后运行您的工作流程。