简单的节点名称对应关系
简单的节点所需4要素:
- CATEGORY
- INPUT_TYPES
- RETURN_TYPES
- FUNCTION
class ImageSelector:
CATEGORY = "example"
@classmethod
def INPUT_TYPES(s):
return { "required": { "images": ("IMAGE",), } }
RETURN_TYPES = ("IMAGE",)
FUNCTION = "choose_image"
编写choose_image函数
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传入多个图片,image.flatten()如将三通道的图片转换成1维数组,torch.mean计算一个张量的平均值将这个元素作为一个标准的Python数值(如float或int)来使用,就可以使用.item()方法,之后转换成一个list列表max找每一个图像列表最大的数,index查找最大的图像列表下标, 选择亮度最高的图像并添加一个批次维度result = images[brightest].unsqueeze(0),尾随逗号对于确保返回元组至关重要。return (result,)
为了让 Comfy 识别新节点,我们需要将目录转换为 Python 模块, 通过添加 ,如下所示image_selector_init_.py:
from .image_selector_node import ImageSelector
NODE_CLASS_MAPPINGS = {
"Image Selector" : ImageSelector,
}
__all__ = ['NODE_CLASS_MAPPINGS']
all:这个变量确保只有 NODE_CLASS_MAPPINGS 被导出,即使包中有更多的变量和函数,基本上每一个节点编写都需要。
添加一些选项
该节点可能有点无聊,因此我们可能会添加一些选项;一个允许您 选择最亮的图像,或者选择最红、最蓝或最绿的图像。编辑 Python 以添加另一个输入, so 看起来像:INPUT_TYPES
@classmethod
def INPUT_TYPES(s):
return { "required": { "images": ("IMAGE",),
"mode": (["brightest", "reddest", "greenest", "bluest"],)} }
然后更新 main 函数。我们将使用一个相当幼稚的 ‘reddest’ 定义,即像素的平均值除以所有三种颜色的平均值。
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,)
修改 init.py 文件
在 image_selector 目录的 init.py 文件中,你需要指定 WEB_DIRECTORY 变量,以便系统知道在哪里查找JavaScript文件。这通过在 all 列表中添加 WEB_DIRECTORY 来实现,这样当模块被导入时,这个变量也会被导出。
# image_selector/__init__.py
WEB_DIRECTORY = "./js"
__all__ = ['NODE_CLASS_MAPPINGS', 'WEB_DIRECTORY']
编写JavaScript代码
在 image_selector.js 文件中,你可以使用提供的框架和API来扩展客户端的功能。以下是一个示例代码,它演示了如何注册一个扩展并设置一个事件监听器来处理来自服务器的消息。
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);
},
})
__init__新加节点好看的展示:
NODE_DISPLAY_NAME_MAPPINGS = {
"Image Selector" : "Image ♾️Selector",
}