ComfyUI - 游览

最新推荐文章于 2025-04-28 09:38:40 发布

AI小译

最新推荐文章于 2025-04-28 09:38:40 发布

阅读量782

点赞数 25

分类专栏： Github精选文章标签： python ai comfyui

本文链接：https://blog.csdn.net/2401_89025022/article/details/144251489

版权

Github精选专栏收录该内容

32 篇文章

订阅专栏

这个页面将逐步指导你创建一个自定义节点的过程，该节点接收一批图像并返回其中的一张。最初，该节点将返回平均颜色最浅的图像；然后我们将扩展它以拥有一系列选择标准，最后添加一些客户端代码。

这个页面假设你对Python或JavaScript的了解非常有限。

在这个教程之后，深入了解服务器端代码, 客户端代码，或客户端-服务器通信。

基本节点

设置

这个自定义节点的所有代码都将放在单个目录中。因此，请首先在您的 ComfyUI 文件夹中找到 custom_nodes 目录，并在其中创建一个新的目录，例如命名为 image_selector。

这个新目录是所有与新的自定义节点相关的代码的基准目录。

Python 框架

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

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

自定义节点是一个Python类，它必须包括以下四件事：CATEGORY，它指定自定义节点在添加新节点菜单中的位置，INPUT_TYPES，它是一个类方法，定义节点将接受哪些输入（稍后有关返回字典的详细信息，请参阅），RETURN_TYPES，它定义节点将产生哪些输出，以及FUNCTION，执行节点时将调用的函数的名称。

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

添加main函数

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

import torch

然后将函数添加到您的类中。图像的数据类型是带有形状的torch.Tensor``[B,H,W,C]，其中B是批处理大小，C是通道数-3，用于RGB。如果我们迭代这样一个张量，我们将得到一系列B形状张量[H,W,C]。.flatten()方法将其转换为一维张量，长度为H*W*C，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）维度，在这种情况下，维度为零，给我们[B,H,W,C]，B=1：单个图像。
在return (result,)中，尾随逗号对于确保返回元组至关重要。

部署节点

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

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"],)} }

然后更新main函数。我们将使用一个相当简单的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,)

调整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);
    },
})