演练

本页将逐步指导您创建自定义节点的过程。

该节点会获取一批图像并返回其中一张图像。最初，该节点将返回平均颜色最浅的图像；然后我们将对其进行扩展，使其具有一系列选择标准，最后添加一些客户端代码。

本页面假设您对 Python 或 Javascript 的了解很少。

基本节点

代码位置（设置、自定义节点的安装）

（开始这里之前，建议先了解ComfyUI的自定义节点的安装方法（手动安装方法））

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

代码内容

Python框架的四个基本内容

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

class ImageSelector:
    CATEGORY = "example"
    @classmethod    
    def INPUT_TYPES(s):
        return { "required":  { "images": ("IMAGE",), } }
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "choose_image" # 这里利用python的反射机制直接通过字符串调用对应的函数名

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

CATEGORY	节点类别	指定自定义节点在添加新节点菜单中的位置
INPUT_TYPES	输入类型	这是一个类方法，定义节点将采取的输入（稍后请参阅返回的字典的详细信息）
RETURN_TYPES	返回类型	定义节点将产生的输出
FUNCTION	函数	执行节点时将调用的函数的名称

Tip

IMAGE 请注意，尽管我们期望接收一批图像并仅返回一张，但输入和输出的数据类型为(单数)。在 Comfy 中，IMAGE 表示图像批次，单个图像被视为大小为 1 的批次。

添加主要函数

这个函数的功能需要有一定的torch代码基础

import torch # 由于我们处理的图像在内部存储为 `torch.Tensor`

...

def choose_image(self, images):
    """
    功能：
    接收 中定义的命名参数 `INPUT_TYPES`
    并返回 `tuple` 中定义的 `RETURN_TYPES`
    """
    brightness = list(torch.mean(image.flatten()).item() for image in images)
    brightest = brightness.index(max(brightness))
    # `images[brightest]`将返回一个形状为的张量 `[H,W,C]`。`unsqueeze` 用于在零维度处插入一个（长度为 1）维度，为我们提供 `[B,H,W,C]`：`B=1`单个图像。
    result = images[brightest].unsqueeze(0)
    return (result,) # 尾随逗号是至关重要的，以确保您返回一个元组

将函数添加到您的类中。图像的数据类型为 torch.Tensor 形状为 [B,H,W,C]。其中：

B 是批处理大小，
C 是通道数，RGB 时为 3
H,W 是宽高
如果我们迭代这样的张量，我们将得到一系列 B 形状为的张量 [H,W,C]。该 .flatten() 方法将其转换为长度为的一维张量 H*W*C，torch.mean() 取平均值，并将 .item() 单值张量转换为 Python 浮点数。

Python模块的形式

为了让 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 服务器，重新启动，重新加载网页，并运行您的工作流程。

链接到当前文件 0

没有文件链接到当前文件