# kijai/comfyui-kwaikolorswrapper architecture

> Architecture diagram of the GitHub repository kijai/comfyui-kwaikolorswrapper, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-17). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/kijai/comfyui-kwaikolorswrapper
- Repository: https://github.com/kijai/comfyui-kwaikolorswrapper

## Overview

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: [`kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py`](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py) (Kolors pipeline), [`kolors/models/tokenization_chatglm.py`](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/tokenization_chatglm.py) (ChatGLM tokenizer), [`nodes.py`](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/nodes.py) (nodes.py), [`kolors/models/configuration_chatglm.py`](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/configuration_chatglm.py) (ChatGLM config)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_integration["ComfyUI integration"]
  node_nodes_py["nodes.py<br/>comfyui node entry<br/>[nodes.py]"]
  node_init_py["__init__.py<br/>package entry<br/>[__init__.py]"]
end

subgraph group_runtime["Kolors runtime"]
  node_pipe_main["Kolors pipeline<br/>diffusers pipeline"]
  node_pipe_pkg["pipelines pkg<br/>pipeline package<br/>[__init__.py]"]
  node_models_pkg["models pkg<br/>model package<br/>[__init__.py]"]
  node_chatglm_cfg["ChatGLM config<br/>model config"]
  node_chatglm_model["ChatGLM model<br/>text encoder"]
  node_chatglm_tok["ChatGLM tokenizer"]
  node_quant{{"Quantization<br/>load/runtime utility<br/>[quantization.py]"}}
end

subgraph group_assets["Bundled assets"]
  node_text_cfg["Encoder config<br/>static config"]
  node_tok_model["Tokenizer model<br/>tokenizer asset<br/>[tokenizer.model]"]
  node_tok_cfg["Tokenizer config<br/>tokenizer asset"]
  node_tok_vocab["Token vocab<br/>tokenizer asset<br/>[vocab.txt]"]
  node_example_t2i["Text2img example<br/>workflow example"]
  node_example_i2i["Img2img example<br/>workflow example"]
end

node_packaging["Packaging<br/>release config<br/>[pyproject.toml]"]
node_deps["Dependencies<br/>install config<br/>[requirements.txt]"]
node_publish["Publish workflow<br/>ci workflow<br/>[publish.yml]"]

node_nodes_py -->|"invokes"| node_pipe_main
node_init_py -->|"registers"| node_nodes_py
node_pipe_main -->|"tokenizes"| node_chatglm_tok
node_pipe_main -->|"encodes"| node_chatglm_model
node_pipe_main -->|"loads"| node_quant
node_pipe_main -->|"configures"| node_chatglm_cfg
node_chatglm_model -->|"uses"| node_chatglm_cfg
node_chatglm_tok -->|"reads"| node_tok_model
node_chatglm_tok -->|"reads"| node_tok_cfg
node_chatglm_tok -->|"reads"| node_tok_vocab
node_chatglm_cfg -->|"bootstraps"| node_text_cfg
node_pipe_main -.->|"shown in"| node_example_t2i
node_pipe_main -.->|"shown in"| node_example_i2i
node_nodes_py -.->|"packaged by"| node_packaging
node_packaging -.->|"installs"| node_deps
node_publish -.->|"publishes"| node_packaging

click node_nodes_py "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/nodes.py"
click node_init_py "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/__init__.py"
click node_pipe_main "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py"
click node_pipe_pkg "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/pipelines/__init__.py"
click node_models_pkg "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/__init__.py"
click node_chatglm_cfg "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/configuration_chatglm.py"
click node_chatglm_model "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/modeling_chatglm.py"
click node_chatglm_tok "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/tokenization_chatglm.py"
click node_quant "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/quantization.py"
click node_text_cfg "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/text_encoder_config.json"
click node_tok_model "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/tokenizer.model"
click node_tok_cfg "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/tokenizer_config.json"
click node_tok_vocab "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/vocab.txt"
click node_example_t2i "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/examples/kolors_example.json"
click node_example_i2i "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/examples/kolors_img2img_example.json"
click node_packaging "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/pyproject.toml"
click node_deps "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/requirements.txt"
click node_publish "https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/.github/workflows/publish.yml"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_nodes_py,node_init_py toneBlue
class node_pipe_main,node_pipe_pkg,node_models_pkg,node_chatglm_cfg,node_chatglm_model,node_chatglm_tok,node_quant toneAmber
class node_text_cfg,node_tok_model,node_tok_cfg,node_tok_vocab,node_example_t2i,node_example_i2i toneMint
class node_packaging,node_deps,node_publish toneNeutral
```

## Components

### ComfyUI integration

Boundary between ComfyUI graph execution and the custom node package

- [nodes.py](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/nodes.py): comfyui node entry · Defines the custom node interface, parameters, and execution hook for ComfyUI · `nodes.py`
- [__init__.py](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/__init__.py): package entry · Package registration entry for the custom node package · `__init__.py`

### Kolors runtime

Pipeline orchestration and model execution

- [Kolors pipeline](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py): diffusers pipeline · Main diffusion orchestration for text-to-image generation · `kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py`
- [pipelines pkg](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/pipelines/__init__.py): pipeline package · Pipeline package namespace and exports · `kolors/pipelines/__init__.py`
- [models pkg](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/__init__.py): model package · Model namespace and exports for ChatGLM components · `kolors/models/__init__.py`
- [ChatGLM config](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/configuration_chatglm.py): model config · Configuration schema for the ChatGLM text encoder · `kolors/models/configuration_chatglm.py`
- [ChatGLM model](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/modeling_chatglm.py): text encoder · ChatGLM-based language model used to produce prompt embeddings · `kolors/models/modeling_chatglm.py`
- [ChatGLM tokenizer](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/tokenization_chatglm.py): tokenizer · Tokenizer implementation for ChatGLM-compatible text inputs · `kolors/models/tokenization_chatglm.py`
- [Quantization](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/kolors/models/quantization.py): load/runtime utility · Quantized loading paths for reduced-memory text encoder variants · `kolors/models/quantization.py`

### Bundled assets

Static configs, tokenizer files, and example workflows

- [Encoder config](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/text_encoder_config.json): static config · Bundled text encoder bootstrap configuration · `configs/text_encoder_config.json`
- [Tokenizer model](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/tokenizer.model): tokenizer asset · Sentencepiece-style tokenizer model used by the bundled tokenizer · `configs/tokenizer/tokenizer.model`
- [Tokenizer config](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/tokenizer_config.json): tokenizer asset · Tokenizer metadata and special-token settings · `configs/tokenizer/tokenizer_config.json`
- [Token vocab](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/configs/tokenizer/vocab.txt): tokenizer asset · Vocabulary file used by the bundled tokenizer · `configs/tokenizer/vocab.txt`
- [Text2img example](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/examples/kolors_example.json): workflow example · ComfyUI workflow example for text-to-image generation · `examples/kolors_example.json`
- [Img2img example](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/examples/kolors_img2img_example.json): workflow example · ComfyUI workflow example for image-to-image generation · `examples/kolors_img2img_example.json`

### Other components

- [Packaging](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/pyproject.toml): release config · Python package metadata and dependency specification · `pyproject.toml`
- [Dependencies](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/requirements.txt): install config · Pinned or declared Python dependencies for the wrapper · `requirements.txt`
- [Publish workflow](https://github.com/kijai/comfyui-kwaikolorswrapper/blob/main/.github/workflows/publish.yml): ci workflow · GitHub Actions workflow for build and publish automation · `.github/workflows/publish.yml`

## Connections

- 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