07-自定义节点与扩展¶

class MyCustomNode:
    """
    一个自定义节点的基本结构
    """

    # 节点在菜单中的分类
    CATEGORY = "my_nodes"

    # 返回类型（输出的数据类型）
    RETURN_TYPES = ("IMAGE",)

    # 返回名称
    RETURN_NAMES = ("image",)

    # 执行函数名
    FUNCTION = "execute"

    @classmethod
    def INPUT_TYPES(cls):
        """定义节点的输入接口"""
        return {
            "required": {
                "image": ("IMAGE",),                    # 图像输入
                "brightness": ("FLOAT", {               # 浮点数参数
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.1,
                    "display": "slider"
                }),
            },
            "optional": {
                "mask": ("MASK",),                      # 可选的遮罩输入
            }
        }

    def execute(self, image, brightness, mask=None):
        """
        节点的执行逻辑
        输入参数与 INPUT_TYPES 中定义的一致
        返回元组，与 RETURN_TYPES 对应
        """
        result = image * brightness
        return (result,)

from .nodes import MyCustomNode

# 节点类映射（内部名称 → 类）
NODE_CLASS_MAPPINGS = {
    "MyCustomNode": MyCustomNode,
}

# 节点显示名称映射（内部名称 → 界面显示名称）
NODE_DISPLAY_NAME_MAPPINGS = {
    "MyCustomNode": "My Custom Node ✨",
}

# ComfyUI/custom_nodes/comfy_brightness_contrast/nodes.py

import torch

class BrightnessContrastNode:
    """
    调节图像的亮度和对比度

    用法：
    - brightness > 1.0 增加亮度，< 1.0 降低亮度
    - contrast > 1.0 增加对比度，< 1.0 降低对比度
    """

    CATEGORY = "image/adjustments"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "adjust"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "brightness": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.05,
                    "display": "slider"
                }),
                "contrast": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.05,
                    "display": "slider"
                }),
            }
        }

    def adjust(self, image: torch.Tensor, brightness: float, contrast: float):
        """
        image: Tensor (B, H, W, C) 范围 [0, 1]
        brightness: 亮度系数
        contrast: 对比度系数
        """
        # 亮度调节
        result = image * brightness

        # 对比度调节（以 0.5 为中心）
        mean = 0.5
        result = (result - mean) * contrast + mean

        # 限制范围到 [0, 1]
        result = torch.clamp(result, 0.0, 1.0)

        return (result,)

class ImageInvertNode:
    """
    反转图像颜色（负片效果）
    """

    CATEGORY = "image/adjustments"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "invert"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "strength": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.1,
                    "display": "slider"
                }),
            }
        }

    def invert(self, image: torch.Tensor, strength: float):
        """
        将图像颜色反转，strength 控制反转强度
        strength = 1.0 完全反转，0.0 不变
        """
        inverted = 1.0 - image
        result = image * (1 - strength) + inverted * strength
        return (result,)

# ComfyUI/custom_nodes/comfy_brightness_contrast/__init__.py

from .nodes import BrightnessContrastNode, ImageInvertNode

NODE_CLASS_MAPPINGS = {
    "BrightnessContrast": BrightnessContrastNode,
    "ImageInvert": ImageInvertNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BrightnessContrast": "Brightness & Contrast 🔆",
    "ImageInvert": "Image Invert 🔄",
}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

@classmethod  # @classmethod类方法，第一个参数为类本身
def INPUT_TYPES(cls):
    return {
        "required": {
            # 下拉选择
            "mode": (["normal", "multiply", "screen", "overlay"],),

            # 整数输入
            "size": ("INT", {
                "default": 512, "min": 64, "max": 4096, "step": 64
            }),

            # 多行文本
            "text": ("STRING", {
                "multiline": True,
                "default": "hello world"
            }),

            # 布尔开关
            "enabled": ("BOOLEAN", {"default": True}),
        }
    }