Klein Pruned Text Encoder Loader

ComfyUI-KleinPrunedTE

ComfyUI loader node for the pruned FLUX.2-klein-9B text encoder (5.1B) from SearchingMan/FLUX.2-klein-9B-Text-Encoder-Pruned-5.1B.

The pruned encoder (5.1B params vs the original 8.2B Qwen3-8B) has mixed per-layer FFN widths and GQA head counts plus remapped hidden-state taps (9, 17, 25) on 26 layers. The stock CLIPLoader builds a uniform Qwen3-8B and fails with size mismatch errors on this checkpoint — this node infers the correct per-layer shapes directly from the checkpoint and wires the remapped taps into ComfyUI's Klein conditioning path. Everything downstream (UNETLoader with the Klein DiT, VAELoader, samplers) is stock.

Why: ~35% smaller text encoder, ~5 GiB less peak VRAM during encoding (10.6 vs 15.5 GiB in bf16), with image output nearly indistinguishable from the full encoder. See the model card for side-by-side comparisons and benchmarks.

Installation

ComfyUI-Manager

Manager → Install via Git URL →

https://github.com/kgonia/ComfyUI-KleinPrunedTE

Manual

cd ComfyUI/custom_nodes
git clone https://github.com/kgonia/ComfyUI-KleinPrunedTE

Restart ComfyUI afterwards. No extra Python dependencies — the node only uses ComfyUI's own APIs.

Requires a ComfyUI version with FLUX.2 support (Flux2Scheduler, EmptyFlux2LatentImage).

Model files

| File | Put in | Source | |------|--------|--------| | flux2-klein-te-pruned-5.1b.safetensors | models/text_encoders/ | SearchingMan/FLUX.2-klein-9B-Text-Encoder-Pruned-5.1B (ComfyUI single-file variant) | | flux-2-klein-9b.safetensors | models/diffusion_models/ | black-forest-labs/FLUX.2-klein-9B (gated — accept the license on Hugging Face and log in to download) | | flux2-vae.safetensors | models/vae/ | Comfy-Org/vae-text-encorder-for-flux-klein-9b |

The example workflow embeds these download links, so ComfyUI will offer to fetch any missing files when you open it (the gated DiT still requires being logged in on huggingface.co).

Usage

1. Open example_workflows/image_flux2_klein_pruned_te_9b.json (Workflow → Open, or drag it onto the canvas).
2. The Load Pruned Klein Text Encoder (5.1B) node replaces CLIPLoader; its CLIP output feeds the usual CLIPTextEncode nodes.
3. Distilled Klein settings: 4 steps, CFG 1.0, euler (already set in the workflow).

Do not load the pruned checkpoint with the stock CLIPLoader — it cannot express the pruned architecture and will fail with size mismatch for model.layers... errors.

How it works

• Reads the checkpoint and infers each layer's FFN width and Q/KV head counts from the tensor shapes (no sidecar config files needed).
• Reads the remapped tap indices from the safetensors metadata key text_encoder_out_layers (falls back to (9, 17, 25)).
• Builds ComfyUI's Klein text-encoder stack with those per-layer shapes, then loads the weights with a strict key check, so a wrong file fails loudly instead of producing garbage.

License

• Code: MIT (this repository).
• Model weights: the pruned text encoder is a derivative of FLUX.2-klein-9B and inherits the FLUX Non-Commercial License. See the weight repositories for details.