ComfyUI自定义节点开发教程用KwaiKolorsWrapper源码学会节点注册、显存调度与offload【免费下载链接】ComfyUI-KwaiKolorsWrapperDiffusers wrapper to run Kwai-Kolors model项目地址: https://gitcode.com/gh_mirrors/co/ComfyUI-KwaiKolorsWrapperComfyUI-KwaiKolorsWrapper 是一个基于 Diffusers 的 ComfyUI 自定义节点用于在 ComfyUI 中运行快手 Kwai-Kolors 文生图模型。本文将带你快速读完它的核心源码掌握 ComfyUI 自定义节点开发的三大核心技能节点注册、显存调度与 offload 显存卸载帮你快速上手开发自己的 ComfyUI 节点。1. 这个节点能做什么ComfyUI 本身不内置所有模型。Kwai-Kolors 是一个中文能力很强的文生图模型它由三部分组成组件说明UNet扩散模型负责去噪生成约 16.5 GBfp16ChatGLM3文本编码器把提示词编码为向量是显存消耗大头SDXL VAE复用 SDXL 的 VAE 解码图像节点提供了 5 个自定义节点覆盖加载模型 → 编码文本 → 采样 → 解码出图的完整链路(Down)load Kolors Model— 自动下载并加载 UNet(Down)load ChatGLM3 Model— 自动下载文本编码器Load ChatGLM3 Model— 从本地 safetensors 加载支持量化模型Kolors Text Encode— 编码正/负向提示词Kolors Sampler— 采样节点支持 6 种采样调度器 官方工作流可直接导入文生图见 examples/kolors_example.json图生图img2img见 examples/kolors_img2img_example.json。2. 快速上手3步安装并跑通 Kolors 节点步骤一克隆节点到 custom_nodes 目录git clone https://gitcode.com/gh_mirrors/co/ComfyUI-KwaiKolorsWrapper将克隆后的文件夹放入 ComfyUI 的custom_nodes目录下。步骤二安装依赖依赖清单在 requirements.txt要求 transformers 版本不低于 4.38.0pip install -r requirements.txt步骤三准备模型模型约 16.5 GB首次运行时会自动下载到ComfyUI/models/diffusers/Kolors。目录结构需包含scheduler、text_encoder、unet三个子目录。显存参考1024×1024 单图文本编码器精度占用显存fp16约 13 GBquant88位量化约 8 GBquant44位量化约 4 GB⚠️ 量化会牺牲部分画质作者经验quant8 效果尚可quant4 画质损失较明显。3. 节点注册机制你的节点如何出现在 ComfyUI 里这是 ComfyUI 节点开发的第一课。核心就两个映射表定义在 nodes.py 文件末尾NODE_CLASS_MAPPINGS节点 ID → 节点类ComfyUI 靠它扫描并注册节点NODE_DISPLAY_NAME_MAPPINGS节点 ID → 显示名称即界面里看到的名字__init__.py只有一行关键代码负责把两个映射表导出给 ComfyUIfrom .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS3.1 一个节点的完整骨架每个节点类需要 4 个类属性以DownloadAndLoadKolorsModel为例属性作用INPUT_TYPES类方法声明输入端口下拉框、数值、字符串等RETURN_TYPES声明输出端口类型RETURN_NAMES输出端口的显示名FUNCTION实际执行的函数名如loadmodelCATEGORY节点在界面菜单中的分类RETURN_TYPES (KOLORSMODEL,) RETURN_NAMES (kolors_model,) FUNCTION loadmodel CATEGORY KwaiKolorsWrapper3.2 自定义数据类型节点间传黑盒对象注意KOLORSMODEL、CHATGLM3MODEL、KOLORS_EMBEDS这些类型——它们是节点自定义的类型名在 ComfyUI 中只是类型标签内部实际是一个字典kolors_model {pipeline: pipeline, dtype: dtype}这种黑盒传递模式非常值得学习加载节点把复杂对象整个 pipeline打包成字典输出下游节点直接使用类型标签保证了端口只能连接到兼容的节点上。4. 显存调度与 offload低显存跑大模型的关键Kolors 的文本编码器高达 13 GBfp16如果 UNet 和文本编码器同时驻留显存普通显卡根本跑不动。作者在 nodes.py 中的 offload 手法是本文最值得抄的作业① 两个设备分工device mm.get_torch_device() # 计算设备GPU offload_device mm.unet_offload_device() # 卸载设备CPU 内存模型平时放在offload_device上休息需要计算时才搬到device。② 用后即走KolorsTextEncode.encode()的显存节奏mm.unload_all_models()mm.soft_empty_cache()— 先把其他模型从显存赶走text_encoder.to(device)— 把文本编码器搬上 GPU完成正向提示词与负向提示词两次编码text_encoder.to(offload_device)— 立刻搬回 CPUmm.soft_empty_cache()gc.collect()— 清空缓存、回收垃圾③ 采样节点同理KolorsSampler.process()中先pipeline.unet.to(device)采样结束后pipeline.unet.to(offload_device)。 核心思路任何时刻显存里只放正在算的模型。编码阶段显存里是 ChatGLM3采样阶段换入 UNet其余时间全部 offload 到 CPU 内存。④ 零内存加载技巧加载 ChatGLM3 时作者用accelerate的init_empty_weights()先空壳创建模型不占内存再配合set_module_tensor_to_device()逐层把权重直接放到目标设备上避免在 CPU 内存中产生一份完整副本。5. 进阶4位/8位量化是怎么省显存的量化实现在 kolors/models/quantization.py核心是一个QuantizedLinear模块它替换掉 ChatGLM3 每层中的 4 个线性层权重压缩fp16 权重按逐行缩放因子压缩为 int84 位时再调用 CUDA 内核把每两个 int8 打包进一个字节前向计算W8A16Linear在计算前把量化权重临时解压回 fp16再执行矩阵乘法触发方式LoadChatGLM3节点根据文件名中的4bit/8bit自动调用text_encoder.quantize(4)或quantize(8)这就是为什么 quant4 能把 13 GB 的文本编码器压进 4 GB 显存。6. 源码导览按这个顺序读效率最高文件内容init.py节点映射表导出入口nodes.py全部 5 个节点注册、加载、编码、采样、offloadkolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py改造自 Diffusers 的 SDXL 采样管线kolors/models/modeling_chatglm.pyChatGLM3 模型结构kolors/models/tokenization_chatglm.pyChatGLM 分词器configs/text_encoder_config.json文本编码器配置本地加载用configs/tokenizer/分词器词表与配置7. 总结从这份源码学到的 4 件事注册靠两个映射表NODE_CLASS_MAPPINGSNODE_DISPLAY_NAME_MAPPINGS配合__init__.py导出即可完成节点注册用自定义类型做数据黑盒复杂对象打包成字典通过类型标签在节点间安全传递offload 是显存管理的核心to(device)与to(offload_device)成对出现配合unload_all_models()、soft_empty_cache()、gc.collect()保持显存干净量化 零内存加载低显存设备的救命组合拳掌握了这套模式你完全可以照着这个骨架把自己想跑的任意 Diffusers 模型包装成 ComfyUI 节点。【免费下载链接】ComfyUI-KwaiKolorsWrapperDiffusers wrapper to run Kwai-Kolors model项目地址: https://gitcode.com/gh_mirrors/co/ComfyUI-KwaiKolorsWrapper创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考