开源 AIGC 工具自定义节点开发教程（以 ComfyUI 为例）

核心概念

1. 节点(Node)：数据处理的基本单元，包含输入/输出端口和计算逻辑
2. 工作流(Workflow)：通过连接节点形成的处理管道
3. 自定义节点：扩展原生功能的核心开发方式

开发流程

1. 环境准备

# 克隆 ComfyUI 源码
git clone https://github.com/comfyanonymous/ComfyUI
cd ComfyUI/custom_nodes

2. 创建节点模板

新建 my_node.py 文件：

import comfy
import torch

class MyCustomNode:
    # 定义输入参数
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_image": ("IMAGE",),
                "strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0})
            }
        }
    
    # 定义输出类型
    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("output_img", "mask")
    
    # 核心处理逻辑
    def process(self, input_image, strength):
        # 示例：简单图像处理
        processed = input_image * strength
        mask = torch.ones_like(input_image[..., 0])
        return (processed, mask)

# 注册节点
comfy.register_node("MyCustomNode", MyCustomNode)

3. 功能扩展模式

4. 模型集成示例

class StableDiffusionXLNode:
    def __init__(self):
        # 延迟加载大模型
        self.model = None
        
    def load_model(self):
        if self.model is None:
            from diffusers import StableDiffusionXLPipeline
            self.model = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
    
    def process(self, prompt, negative_prompt):
        self.load_model()
        result = self.model(prompt, negative_prompt=negative_prompt)
        return (result.images[0],)

5. 调试技巧

1. 日志输出：

   comfy.logger.info(f"Processing image shape: {input_image.shape}")
   

2. 热重载：
   • 修改代码后保存，在 ComfyUI 界面按 Ctrl+Shift+R 刷新节点
3. 单元测试：

   def test_node():
       dummy_input = torch.rand(1, 512, 512, 3)
       node = MyCustomNode()
       out, mask = node.process(dummy_input, 0.7)
       assert out.shape == dummy_input.shape
   

最佳实践

1. 性能优化：
   • 大模型使用 @lru_cache 缓存
   • 避免在 process() 中重复初始化
2. 错误处理：

   try:
       return self.model.generate(...)
   except RuntimeError as e:
       comfy.logger.error(f"VRAM不足: {str(e)}")
       return torch.zeros_like(input_image)
   

3. UI 友好设计：
   • 使用 SLIDER 替代自由输入
   • 添加预览图输出端口

部署流程

1. 将节点文件放入 ComfyUI/custom_nodes 目录
2. 启动 ComfyUI 后节点自动加载
3. 在节点菜单的 Custom 分类中找到新节点

提示：复杂节点建议使用 __init__.py 组织多文件模块，参考官方示例节点结构。开发完成后可打包发布到 ComfyUI 社区商店。

(图：自定义节点开发流程示意图)

通过此教程，您可快速实现以下扩展：

• 集成新型文生图模型（如 SDXL-Lightning）
• 添加视频处理节点
• 创建多模态控制节点
• 开发领域特定优化模块