192 lines
7.5 KiB
Markdown
192 lines
7.5 KiB
Markdown
# Getting Started - Custom Node Walkthrough
|
|
|
|
Source: https://docs.comfy.org/custom-nodes/walkthrough
|
|
|
|
This page will take you step-by-step through the process of creating a custom node.
|
|
|
|
Our example will take a batch of images, and return one of the images. Initially, the node will return the image which is, on average, the lightest in color; we'll then extend it to have a range of selection criteria, and then finally add some client side code.
|
|
|
|
This page assumes very little knowledge of Python or Javascript.
|
|
|
|
After this walkthrough, dive into the details of backend code and frontend code.
|
|
|
|
## Write a basic node
|
|
|
|
### Prerequisites
|
|
|
|
- A working ComfyUI installation. For development, we recommend installing ComfyUI manually.
|
|
- A working comfy-cli installation.
|
|
|
|
### Setting up
|
|
|
|
```bash
|
|
cd ComfyUI/custom_nodes
|
|
comfy node scaffold
|
|
```
|
|
|
|
After answering a few questions, you'll have a new directory set up.
|
|
|
|
```bash
|
|
~ % 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!
|
|
```
|
|
|
|
### Defining the node
|
|
|
|
Add the following code to the end of `src/nodes.py`:
|
|
|
|
```python
|
|
class ImageSelector:
|
|
CATEGORY = "example"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": { "images": ("IMAGE",), } }
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "choose_image"
|
|
```
|
|
|
|
> **Info**: The basic structure of a custom node is described in detail in the server overview.
|
|
|
|
A custom node is defined using a Python class, which must include these four things:
|
|
- `CATEGORY`: specifies where in the add new node menu the custom node will be located
|
|
- `INPUT_TYPES`: a class method defining what inputs the node will take
|
|
- `RETURN_TYPES`: defines what outputs the node will produce
|
|
- `FUNCTION`: the name of the function that will be called when the node is executed
|
|
|
|
> **Tip**: 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`:
|
|
|
|
```python
|
|
import torch
|
|
```
|
|
|
|
Then add the function to your class. The datatype for image is `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.
|
|
|
|
```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,)
|
|
```
|
|
|
|
Notes on those last two lines:
|
|
- `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]` with `B=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 the `NODE_CLASS_MAPPINGS` variable at the end of `src/nodes.py`. You must restart ComfyUI to see any changes.
|
|
|
|
```python
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Example" : Example,
|
|
"Image Selector" : ImageSelector,
|
|
}
|
|
|
|
# Optionally, you can rename the node in the `NODE_DISPLAY_NAME_MAPPINGS` dictionary.
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Example": "Example Node",
|
|
"Image Selector": "Image Selector",
|
|
}
|
|
```
|
|
|
|
> **Info**: 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 your `INPUT_TYPES` to look like:
|
|
|
|
```python
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": { "images": ("IMAGE",),
|
|
"mode": (["brightest", "reddest", "greenest", "bluest"],)} }
|
|
```
|
|
|
|
Then update the main function. We'll use a fairly naive definition of 'reddest' as being the average `R` value of the pixels divided by the average of all three colors. So:
|
|
|
|
```python
|
|
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,)
|
|
```
|
|
|
|
## 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:
|
|
|
|
```python
|
|
from server import PromptServer
|
|
```
|
|
|
|
and, at the end of the `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)
|
|
|
|
```python
|
|
PromptServer.instance.send_sync("example.imageselector.textmessage", {"message":f"Picked image {best+1}"})
|
|
return (result,)
|
|
```
|
|
|
|
### 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`:
|
|
|
|
```python
|
|
WEB_DIRECTORY = "./web/js"
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'WEB_DIRECTORY']
|
|
```
|
|
|
|
The client extension is saved as a `.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).
|
|
|
|
```javascript
|
|
app.registerExtension({
|
|
name: "example.imageselector",
|
|
async setup() {
|
|
function messageHandler(event) { alert(event.detail.message); }
|
|
app.api.addEventListener("example.imageselector.textmessage", messageHandler);
|
|
},
|
|
})
|
|
```
|
|
|
|
All we've done is register an extension and add a listener for the message type we are sending in the `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 in the ComfyUI documentation. You can download the example workflow JSON file or view it in the documentation.
|