Write a basic node
Prerequisites
- A working ComfyUI installation. For development, we recommend installing ComfyUI manually.
- A working comfy-cli installation.
Setting up
Defining the node
Add the following code to the end ofsrc/nodes.py
:
src/nodes.py
The basic structure of a custom node is described in detail here.
CATEGORY
,
which specifies where in the add new node menu the custom node will be located,
INPUT_TYPES
, which is a class method defining what inputs the node will take
(see later for details of the dictionary returned),
RETURN_TYPES
, which defines what outputs the node will produce, and FUNCTION
, the name
of the function that will be called when the node is executed.
Notice that the data type for input and output is
IMAGE
(singular) even though
we expect to receive a batch of images, and return just one. In Comfy, IMAGE
means
image batch, and a single image is treated as a batch of size 1.The main function
The main function,choose_image
, receives named arguments as defined in INPUT_TYPES
, and
returns a tuple
as defined in RETURN_TYPES
. Since we’re dealing with images, which are internally
stored as torch.Tensor
,
torch.Tensor
with shape [B,H,W,C]
,
where B
is the batch size and C
is the number of channels - 3, for RGB. If we iterate over such
a tensor, we will get a series of B
tensors of shape [H,W,C]
. The .flatten()
method turns
this into a one dimensional tensor, of length H*W*C
, torch.mean()
takes the mean, and .item()
turns a single value tensor into a Python float.
images[brightest]
will return a Tensor of shape[H,W,C]
.unsqueeze
is used to insert a (length 1) dimension at, in this case, dimension zero, to give us[B,H,W,C]
withB=1
: a single image.- in
return (result,)
, the trailing comma is essential to ensure you return a tuple.
Register the node
To make Comfy recognize the new node, it must be available at the package level. Modify theNODE_CLASS_MAPPINGS
variable at the end of src/nodes.py
. You must restart ComfyUI to see any changes.
src/nodes.py
For a detailed explanation of how ComfyUI discovers and loads custom nodes, see the node lifecycle documentation.
Add some options
That node is maybe a bit boring, so we might add some options; a widget that allows you to choose the brightest image, or the reddest, bluest, or greenest. Edit yourINPUT_TYPES
to look like:
R
value of the pixels divided by the average of all three colors. So:
Tweak the UI
Maybe we’d like a bit of visual feedback, so let’s send a little text message to be displayed.Send a message from server
This requires two lines to be added to the Python code:choose_image
method, add a line to send a message to the front end (send_sync
takes a message
type, which should be unique, and a dictionary)
Write a client extension
To add some Javascript to the client, create a subdirectory,web/js
in your custom node directory, and modify the end of __init__.py
to tell Comfy about it by exporting WEB_DIRECTORY
:
.js
file in the web/js
subdirectory, so create image_selector/web/js/imageSelector.js
with the
code below. (For more, see client side coding).
setup()
method. This reads the dictionary we sent (which is stored in event.detail
).
Stop the Comfy server, start it again, reload the webpage, and run your workflow.
The complete example
The complete example is available here. You can download the example workflow JSON file or view it below: