Crear nodos personalizados en ComfyUI

Esta página te guiará paso a paso por el proceso de crear un nodo personalizado. Nuestro ejemplo tomará un lote de imágenes y devolverá una de ellas. Inicialmente, el nodo devolverá la imagen que, en promedio, sea la más clara de color; luego lo extenderemos para que tenga una serie de criterios de selección y, finalmente, agregaremos algo de código del lado del cliente. Esta página asume muy poco conocimiento de Python o Javascript. Después de este recorrido, profundiza en los detalles del código del backend y el código del frontend.

Escribir un nodo básico

Requisitos previos

Una instalación de ComfyUI en funcionamiento. Para desarrollo, recomendamos instalar ComfyUI manualmente.
Una instalación de comfy-cli en funcionamiento.

Configurar el entorno

cd ComfyUI/custom_nodes
comfy node scaffold

Después de responder algunas preguntas, tendrás un nuevo directorio configurado.

 ~  % comfy node scaffold
You've downloaded .cookiecutters/cookiecutter-comfy-extension before. Is it okay to delete and re-download it? [y/n] (y): y
  [1/9] full_name (): Comfy
  [2/9] email (you@gmail.com): me@comfy.org
  [3/9] github_username (your_github_username): comfy
  [4/9] project_name (My Custom Nodepack): FirstComfyNode
  [5/9] project_slug (firstcomfynode): 
  [6/9] project_short_description (A collection of custom nodes for ComfyUI): 
  [7/9] version (0.0.1): 
  [8/9] Select open_source_license
    1 - GNU General Public License v3
    2 - MIT license
    3 - BSD license
    4 - ISC license
    5 - Apache Software License 2.0
    6 - Not open source
    Choose from [1/2/3/4/5/6] (1): 1
  [9/9] include_web_directory_for_custom_javascript [y/n] (n): y
Initialized empty Git repository in firstcomfynode/.git/
✓ Custom node project created successfully!

Definir el nodo

Agrega el siguiente código al final de src/nodes.py:

src/nodes.py

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

La estructura básica de un nodo personalizado se describe en detalle aquí.

Un nodo personalizado se define con una clase de Python, que debe incluir estas cuatro cosas: CATEGORY, que especifica en qué parte del menú de agregar nuevo nodo se ubicará el nodo personalizado, INPUT_TYPES, que es un método de clase que define qué entradas recibirá el nodo (consulta más adelante los detalles del diccionario que devuelve), RETURN_TYPES, que define qué salidas producirá el nodo, y FUNCTION, el nombre de la función que se llamará cuando se ejecute el nodo.

Observa que el tipo de dato de entrada y salida es IMAGE (singular) aunque esperamos recibir un lote de imágenes y devolver solo una. En Comfy, IMAGE significa lote de imágenes, y una imagen individual se considera un lote de tamaño 1.

La función principal

La función principal, choose_image, recibe argumentos con nombre tal como se definen en INPUT_TYPES y devuelve una tuple tal como se define en RETURN_TYPES. Como trabajamos con imágenes, que internamente se almacenan como torch.Tensor,

import torch

Luego agrega la función a tu clase. El tipo de dato de la imagen es torch.Tensor con forma [B,H,W,C], donde B es el tamaño del lote y C es el número de canales: 3, para RGB. Si iteramos sobre ese tensor, obtendremos una serie de B tensores con forma [H,W,C]. El método .flatten() lo convierte en un tensor unidimensional de longitud H*W*C, torch.mean() toma la media y .item() convierte un tensor de un solo valor en un float de 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,)

Notas sobre esas dos últimas líneas:

images[brightest] devolverá un tensor con forma [H,W,C]. unsqueeze inserta una dimensión (de longitud 1), en este caso, en la dimensión cero, para darnos [B,H,W,C] con B=1: una sola imagen.
en return (result,), la coma final es esencial para garantizar que devuelves una tupla.

Registrar el nodo

Para que Comfy reconozca el nuevo nodo, debe estar disponible a nivel del paquete. Modifica la variable NODE_CLASS_MAPPINGS al final de src/nodes.py. Debes reiniciar ComfyUI para ver los cambios.

src/nodes.py

NODE_CLASS_MAPPINGS = {
    "Example" : Example,
    "Image Selector" : ImageSelector,
}

# Opcionalmente, puedes renombrar el nodo en el diccionario `NODE_DISPLAY_NAME_MAPPINGS`.
NODE_DISPLAY_NAME_MAPPINGS = {
    "Example": "Example Node",
    "Image Selector": "Image Selector",
}

Para una explicación detallada de cómo ComfyUI descubre y carga los nodos personalizados, consulta la documentación del ciclo de vida de los nodos.

Agregar algunas opciones

Ese nodo quizá sea un poco aburrido, así que podríamos agregar algunas opciones: un widget que te permita elegir la imagen más brillante, o la más roja, la más azul o la más verde. Edita tu INPUT_TYPES para que se vea así:

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

Luego actualiza la función principal. Usaremos una definición bastante ingenua de “la más roja” como el valor promedio de R de los píxeles dividido entre el promedio de los tres colores. Así:

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,)

Ajustar la interfaz

Tal vez queramos un poco de retroalimentación visual, así que enviemos un pequeño mensaje de texto para mostrarlo.

Enviar un mensaje desde el servidor

Esto requiere agregar dos líneas al código de Python:

from server import PromptServer

y, al final del método choose_image, agrega una línea para enviar un mensaje al frontend (send_sync recibe un tipo de mensaje, que debe ser único, y un diccionario)

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

Escribir una extensión de cliente

Para agregar un poco de JavaScript al cliente, crea un subdirectorio, web/js, en el directorio de tu nodo personalizado y modifica el final de __init__.py para avisarle a Comfy sobre ello exportando WEB_DIRECTORY:

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

La extensión de cliente se guarda como un archivo .js en el subdirectorio web/js, así que crea image_selector/web/js/imageSelector.js con el código de abajo. (Para más información, consulta programación del lado del cliente).

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

Lo único que hemos hecho es registrar una extensión y agregar un receptor de eventos para el tipo de mensaje que enviamos en el método setup(). Esto lee el diccionario que enviamos (que se almacena en event.detail). Detén el servidor de Comfy, inícialo de nuevo, recarga la página web y ejecuta tu flujo de trabajo.

El ejemplo completo

El ejemplo completo está disponible aquí. Puedes descargar el archivo JSON del flujo de trabajo de ejemplo o verlo a continuación:

Flujo de trabajo del selector de imágenes

Desarrollar con Comfy

Desplegar ComfyUI

Comfy CLI

Specifications

Nodos personalizados

Crear nodos personalizados en ComfyUI

Escribir un nodo básico

Requisitos previos

Configurar el entorno

Definir el nodo

La función principal

Registrar el nodo

Agregar algunas opciones

Ajustar la interfaz

Enviar un mensaje desde el servidor

Escribir una extensión de cliente

El ejemplo completo