How kijai/comfyui-kwaikolorswrapper works
1. What this repository is
- This is a ComfyUI custom node package that wraps the Kwai-Kolors text-to-image pipeline.
- It does not implement a standalone app; it plugs into ComfyUI’s node graph and delegates image generation to a Diffusers-based pipeline.
2. Top-level integration boundary
- The main entry point for ComfyUI integration is nodes.py.
- That file is typically where the custom node classes, parameter schema, and execution hooks live, so it is the boundary between ComfyUI graph inputs and the internal model/pipeline code.
3. Generation pipeline subsystem
- The core image generation logic lives under kolors/pipelines/, especially:
- pipeline_stable_diffusion_xl_chatglm_256.py
- This indicates the pipeline is structurally similar to Stable Diffusion XL, but replaces the text-conditioning path with a ChatGLM-based encoder.
- The pipeline is the main orchestration layer that combines text conditioning, UNet denoising, scheduler steps, and VAE decoding.
4. Text encoder / language model subsystem
- The text encoder implementation is under kolors/models/, specifically:
- configuration_chatglm.py
- modeling_chatglm.py
- tokenization_chatglm.py
- quantization.py
- This is the main non-standard component in the repo: it provides the ChatGLM3-compatible text encoder, tokenizer handling, and quantized loading paths.
- The README highlights that the text encoder is the major VRAM consumer and can be loaded in fp16, 8-bit, or 4-bit variants.
5. Model configuration and tokenizer assets
- The repo includes local configuration files under configs/, notably:
- configs/text_encoder_config.json
- configs/tokenizer/*
- These bundled files let the wrapper load a single safetensors ChatGLM3 model or repo-local tokenizer/config without depending entirely on remote model metadata.
- In architecture terms, these files are “static model descriptors” used during model bootstrap.
6. External model/runtime dependencies
- The generation stack depends on Diffusers, Transformers, and the ComfyUI runtime.
- The README states the Kolors base model is downloaded from Hugging Face into ComfyUI/models/diffusers/Kolors.
- The text encoder can also be loaded from a Hugging Face-hosted ChatGLM3 safetensors checkpoint placed in ComfyUI/models/LLM/checkpoints.
7. Model composition and data flow
- The main inference flow is:
1. ComfyUI passes prompt/settings into the custom node.
2. The node loads or reuses the Kolors pipeline.
3. The tokenizer and ChatGLM text encoder convert the prompt into embeddings.
4. The Diffusers UNet denoises latent noise using those embeddings.
5. The scheduler controls sampling steps.
6. The VAE decodes latents into the final image.
- The README notes that the base SDXL VAE is used, so the image decoding boundary is external to the custom text encoder code.
8. Model layout expectations
- The repo assumes a specific Diffusers model directory layout under ComfyUI/models/diffusers/Kolors.
- That folder contains the model index, scheduler config, text encoder weights/config/tokenizer, and the UNet weights.
- This is important for the architecture diagram because model loading is filesystem-convention driven, not dynamically discovered.
9. Examples and usage artifacts
- The examples/ directory contains JSON example workflows:
- kolors_example.json
- kolors_img2img_example.json
- These are ComfyUI workflow examples showing how the node plugs into a graph for text-to-image and image-to-image style use cases.
10. Packaging and release tooling
- The repository includes pyproject.toml, requirements.txt, and a GitHub Actions workflow in .github/workflows/publish.yml.
- These support Python packaging/dependency installation and automated publishing, but they do not change the inference architecture.
11. Key architectural distinction
- The repo cleanly separates:
- UI/runtime integration: nodes.py
- Pipeline orchestration: kolors/pipelines/
- Model/tokenizer implementation: kolors/models/
- Static configs/assets: configs/
- User-facing examples: examples/
- For an architecture diagram, the important flow is ComfyUI graph → custom node → Kolors pipeline → ChatGLM text encoder + SDXL-like diffusion stack → image output.
Main entry points: pipeline_stable_diffusion_xl_chatglm_256.py Kolors pipeline · tokenization_chatglm.py ChatGLM tokenizer · nodes.py nodes.py · configuration_chatglm.py ChatGLM config
Components 18
ComfyUI integration
Boundary between ComfyUI graph execution and the custom node package
- nodes.py · Defines the custom node interface, parameters, and execution hook for ComfyUI
nodes.py - __init__.py · Package registration entry for the custom node package
__init__.py
Kolors runtime
Pipeline orchestration and model execution
- Kolors pipeline · Main diffusion orchestration for text-to-image generation
kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py - pipelines pkg · Pipeline package namespace and exports
kolors/pipelines/__init__.py - models pkg · Model namespace and exports for ChatGLM components
kolors/models/__init__.py - ChatGLM config · Configuration schema for the ChatGLM text encoder
kolors/models/configuration_chatglm.py - ChatGLM model · ChatGLM-based language model used to produce prompt embeddings
kolors/models/modeling_chatglm.py - ChatGLM tokenizer · Tokenizer implementation for ChatGLM-compatible text inputs
kolors/models/tokenization_chatglm.py - Quantization · Quantized loading paths for reduced-memory text encoder variants
kolors/models/quantization.py
Bundled assets
Static configs, tokenizer files, and example workflows
- Encoder config · Bundled text encoder bootstrap configuration
configs/text_encoder_config.json - Tokenizer model · Sentencepiece-style tokenizer model used by the bundled tokenizer
configs/tokenizer/tokenizer.model - Tokenizer config · Tokenizer metadata and special-token settings
configs/tokenizer/tokenizer_config.json - Token vocab · Vocabulary file used by the bundled tokenizer
configs/tokenizer/vocab.txt - Text2img example · ComfyUI workflow example for text-to-image generation
examples/kolors_example.json - Img2img example · ComfyUI workflow example for image-to-image generation
examples/kolors_img2img_example.json
Other components
- Packaging · Python package metadata and dependency specification
pyproject.toml - Dependencies · Pinned or declared Python dependencies for the wrapper
requirements.txt - Publish workflow · GitHub Actions workflow for build and publish automation
.github/workflows/publish.yml
Connections 16
- nodes.py → Kolors pipeline: invokes · Custom node hands ComfyUI inputs to the Kolors diffusion pipeline
- __init__.py → nodes.py: registers · Package entry exposes the custom nodes to ComfyUI
- Kolors pipeline → ChatGLM tokenizer: tokenizes · Pipeline uses the tokenizer to prepare prompt text
- Kolors pipeline → ChatGLM model: encodes · Pipeline requests text embeddings from the ChatGLM encoder
- Kolors pipeline → Quantization: loads · Pipeline can initialize the encoder through quantized loading paths
- Kolors pipeline → ChatGLM config: configures · Pipeline relies on ChatGLM configuration when constructing the text encoder
- ChatGLM model → ChatGLM config: uses · Model implementation follows the ChatGLM configuration schema
- ChatGLM tokenizer → Tokenizer model: reads · Tokenizer implementation consumes the bundled tokenizer model asset
- ChatGLM tokenizer → Tokenizer config: reads · Tokenizer implementation consumes tokenizer metadata
- ChatGLM tokenizer → Token vocab: reads · Tokenizer implementation consumes the bundled vocabulary file
- ChatGLM config → Encoder config: bootstraps · Configuration can be seeded from the bundled encoder config
- Kolors pipeline → Text2img example: shown in · Example workflow demonstrates text-to-image usage
- Kolors pipeline → Img2img example: shown in · Example workflow demonstrates image-to-image usage
- nodes.py → Packaging: packaged by · Node package is distributed through standard Python packaging metadata
- Packaging → Dependencies: installs · Packaging references declared dependencies
- Publish workflow → Packaging: publishes · CI workflow automates release of the packaged project