ComfyUI/HOW_TO_WRITE_NODES.md

How to Write Nodes in ComfyUI-Inference-Scaling

This guide will walk you through creating custom nodes for ComfyUI using the modern API system.

Overview

In this ComfyUI project, nodes are created using the ComfyAPI system. Each node is a class that:

1. Inherits from IO.ComfyNode
2. Defines its schema (inputs, outputs, metadata) via define_schema()
3. Implements its execution logic via execute()
4. Is registered through a ComfyExtension class

Basic Structure

1. Node Class

Every node is a class that inherits from IO.ComfyNode:

from comfy_api.latest import IO, ComfyExtension
from typing_extensions import override
from inspect import cleandoc

class MyCustomNode(IO.ComfyNode):
    """
    Description of what your node does.
    This docstring will appear in the UI.
    """

    @classmethod
    def define_schema(cls):
        return IO.Schema(
            node_id="MyCustomNode",
            display_name="My Custom Node",
            category="mycategory/subcategory",
            description=cleandoc(cls.__doc__ or ""),
            inputs=[
                # Define inputs here
            ],
            outputs=[
                # Define outputs here
            ],
        )

    @classmethod
    async def execute(cls, ...) -> IO.NodeOutput:
        # Implementation here
        pass

2. Schema Definition

The define_schema() method defines:

• node_id: Unique identifier (usually matches class name)
• display_name: Name shown in the UI
• category: Where it appears in the node menu (use / for subcategories)
• description: Tooltip/help text
• inputs: List of input definitions
• outputs: List of output definitions
• hidden: Optional hidden inputs (like auth tokens)
• is_api_node: Set to True for API nodes

3. Input Types

Common input types available:

# String input
IO.String.Input(
    "prompt",
    default="",
    multiline=True,  # For longer text
    tooltip="Help text shown on hover",
    optional=True,  # Makes it optional
)

# Integer input
IO.Int.Input(
    "seed",
    default=0,
    min=0,
    max=100,
    step=1,
    display_mode=IO.NumberDisplay.slider,  # or .number
    control_after_generate=True,  # Shows randomize button
    tooltip="Random seed",
)

# Float input
IO.Float.Input(
    "strength",
    default=0.5,
    min=0.0,
    max=1.0,
    step=0.01,
    tooltip="Strength value",
)

# Combo/Dropdown input
IO.Combo.Input(
    "model",
    options=["option1", "option2", "option3"],
    default="option1",
    tooltip="Select a model",
)

# Image input
IO.Image.Input(
    "image",
    tooltip="Input image",
    optional=True,
)

# Mask input
IO.Mask.Input(
    "mask",
    tooltip="Mask for inpainting",
    optional=True,
)

# Audio input
IO.Audio.Input(
    "audio",
    tooltip="Input audio",
    optional=True,
)

# Video input
IO.Video.Input(
    "video",
    tooltip="Input video",
    optional=True,
)

4. Output Types

Common output types:

# Image output
IO.Image.Output()

# Audio output
IO.Audio.Output()

# Video output
IO.Video.Output()

# String output
IO.String.Output()

# Integer output
IO.Int.Output()

# Float output
IO.Float.Output()

5. Execute Method

The execute() method is where your node's logic runs:

@classmethod
async def execute(
    cls,
    # Parameters match input names from define_schema
    prompt: str,
    seed: int = 0,
    image: Optional[torch.Tensor] = None,
    # ... other inputs
) -> IO.NodeOutput:
    """
    Execute the node logic.
    
    Args:
        prompt: Text prompt
        seed: Random seed
        image: Optional image tensor (shape: [B, H, W, C])
        ...
    
    Returns:
        IO.NodeOutput with the result
    """
    # Your implementation here
    
    # For image outputs:
    result_tensor = ...  # torch.Tensor with shape [B, H, W, C]
    return IO.NodeOutput(result_tensor)
    
    # For multiple outputs:
    return IO.NodeOutput(image=result_image, metadata=result_metadata)

Important Notes:

• The method is async - use await for async operations
• Input parameters match the names from define_schema()
• Image tensors have shape [Batch, Height, Width, Channels] (usually RGBA)
• Use IO.NodeOutput() to return results

6. Extension Registration

To register your nodes, create an extension class:

class MyExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
        return [
            MyCustomNode,
            AnotherNode,
            # ... list all your nodes
        ]


async def comfy_entrypoint() -> MyExtension:
    """
    Entry point function that ComfyUI calls to load your extension.
    """
    return MyExtension()

Complete Example

Here's a complete example of a simple image processing node:

from io import BytesIO
import torch
import numpy as np
from PIL import Image
from typing import Optional
from typing_extensions import override
from inspect import cleandoc

from comfy_api.latest import IO, ComfyExtension
from comfy_api_nodes.util import validate_string


class ImageBrightnessNode(IO.ComfyNode):
    """
    Adjusts the brightness of an input image.
    """

    @classmethod
    def define_schema(cls):
        return IO.Schema(
            node_id="ImageBrightnessNode",
            display_name="Image Brightness",
            category="image/processing",
            description=cleandoc(cls.__doc__ or ""),
            inputs=[
                IO.Image.Input(
                    "image",
                    tooltip="Input image to adjust",
                ),
                IO.Float.Input(
                    "brightness",
                    default=1.0,
                    min=0.0,
                    max=2.0,
                    step=0.1,
                    tooltip="Brightness multiplier (1.0 = no change)",
                ),
            ],
            outputs=[
                IO.Image.Output(),
            ],
        )

    @classmethod
    async def execute(
        cls,
        image: torch.Tensor,
        brightness: float = 1.0,
    ) -> IO.NodeOutput:
        """
        Adjust image brightness.
        
        Args:
            image: Input image tensor [B, H, W, C]
            brightness: Brightness multiplier
            
        Returns:
            Brightness-adjusted image
        """
        # Ensure we have a batch dimension
        if len(image.shape) == 3:
            image = image.unsqueeze(0)
        
        # Apply brightness adjustment
        # Clamp values to [0, 1] range
        adjusted = torch.clamp(image * brightness, 0.0, 1.0)
        
        return IO.NodeOutput(adjusted)


class MyExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
        return [
            ImageBrightnessNode,
        ]


async def comfy_entrypoint() -> MyExtension:
    return MyExtension()

API Node Example

For API nodes (nodes that call external APIs), you'll typically:

1. Use utility functions from comfy_api_nodes.util:

   • sync_op() - for synchronous API calls
   • poll_op() - for polling async operations
   • validate_string() - for input validation
   • tensor_to_bytesio() - convert image tensors to bytes
   • bytesio_to_image_tensor() - convert bytes to image tensors

2. Include hidden inputs for authentication:

hidden=[
    IO.Hidden.auth_token_comfy_org,
    IO.Hidden.api_key_comfy_org,
    IO.Hidden.unique_id,
],
is_api_node=True,

3. Use ApiEndpoint for API calls:

from comfy_api_nodes.util import ApiEndpoint, sync_op

response = await sync_op(
    cls,
    ApiEndpoint(path="/api/endpoint", method="POST"),
    response_model=YourResponseModel,
    data=YourRequestModel(...),
    files={...},  # Optional file uploads
    content_type="application/json",
)

File Organization

1. Create your node file: comfy_api_nodes/nodes_yourname.py
2. Follow naming conventions:
   • Node classes: YourNodeName (PascalCase)
   • Extension class: YourExtension (PascalCase)
   • File: nodes_yourname.py (snake_case)
3. Import required modules:
   • from comfy_api.latest import IO, ComfyExtension
   • from typing_extensions import override
   • from inspect import cleandoc

Testing Your Node

1. Start ComfyUI:

   python main.py
   

2. Check the node appears in the node menu under your specified category

3. Test the node by:

   • Adding it to a workflow
   • Connecting inputs
   • Executing the workflow

Common Patterns

Working with Images

# Image tensor shape: [Batch, Height, Width, Channels]
# Channels are usually RGBA (4 channels)

# Convert PIL Image to tensor
pil_img = Image.open("image.png").convert("RGBA")
arr = np.asarray(pil_img).astype(np.float32) / 255.0
tensor = torch.from_numpy(arr).unsqueeze(0)  # Add batch dimension

# Convert tensor to PIL Image
tensor = image.squeeze(0).cpu()  # Remove batch, move to CPU
image_np = (tensor.numpy() * 255).astype(np.uint8)
pil_img = Image.fromarray(image_np)

Validation

from comfy_api_nodes.util import validate_string

# Validate string inputs
validate_string(prompt, strip_whitespace=False)

# Check optional inputs
if image is not None:
    # Process image
    pass

Error Handling

if some_condition:
    raise Exception("Error message here")

Tips

1. Use type hints - They help with IDE autocomplete and documentation
2. Add tooltips - Help users understand what each input does
3. Use cleandoc() - Cleans up docstrings for display
4. Make inputs optional when appropriate - Improves usability
5. Follow existing patterns - Look at nodes_openai.py or nodes_stability.py for examples
6. Test thoroughly - Especially edge cases (None inputs, empty strings, etc.)

Resources

• Existing node examples: comfy_api_nodes/nodes_*.py
• ComfyAPI documentation: comfy_api/latest/
• Utility functions: comfy_api_nodes/util/

Next Steps

1. Look at existing nodes for reference
2. Start with a simple node
3. Test incrementally
4. Add more features as needed

Happy node writing! 🎨