Load Pruned Klein Text Encoder (5.1B)

The pruned Klein text encoder exists because the full one is a quarter of your problem. FLUX.2-klein-9B ships with a Qwen3-8B encoder - 8.2B parameters, 14 GiB+ of weights, and ~15.5 GiB of peak VRAM during encoding in bf16. On a 16 GB card that's the difference between the whole fp8 pipeline fitting and the model shuffling in and out of RAM. This node loads the community-pruned 5.1B version (38% smaller) that the stock CLIPLoader physically cannot.

Why can't stock ComfyUI load it? The pruned checkpoint from SearchingMan/FLUX.2-klein-9B-Text-Encoder-Pruned-5.1B doesn't have a uniform architecture anymore. It ships with mixed per-layer FFN widths and GQA head counts, plus the hidden-state taps remapped from (9, 18, 27) to (9, 17, 25) because two layers were deleted. CLIPLoader builds a fixed uniform Qwen3-8B, so it dies with size mismatch for model.layers... errors. KleinPrunedTELoader exists to load that file instead.

How it works

This is one of those rare nodes that reads the checkpoint to figure out its own shape. On load it:

• scans every model.layers.N.* tensor and infers each layer's FFN width and Q/KV head count straight from the shapes - no sidecar config files,
• reads the remapped tap indices from the safetensors metadata key text_encoder_out_layers (falling back to (9, 17, 25)),
• builds ComfyUI's Klein text-encoder stack with those per-layer shapes and loads the weights with a strict key check, so pointing it at the wrong file fails loudly instead of quietly producing garbage.

The author (kgonia, who posts as ThaJedi) did the pruning himself: deleted two layers, SLERP-merged their neighbours, pruned FFN width from 12288 → 8192 with activation-aware Wanda pruning, trimmed GQA KV groups per layer via a sensitivity probe, and ran recovery distillation after each stage. Embedding fidelity against the full encoder is 0.9755 - and in practice, the community reports it as "works just as before." One caveat worth knowing: it's still pruning, so on heavily-iterated prompts where you've memorized the full encoder's exact behaviour you might notice a hair less. Most people won't.

The one input, the one output

The whole node is two sockets:

• clip_name - a dropdown of the .safetensors files in models/text_encoders/. Pick flux2-klein-te-pruned-5.1b.safetensors.
• CLIP output - wire it into your CLIPTextEncode nodes exactly like the output of a CLIPLoader. Everything downstream (the Klein DiT, VAE, samplers) is stock ComfyUI.

That's it. You swap one loader node and keep the rest of your Klein workflow untouched.

Install

Zero Python dependencies - the node only uses ComfyUI's own APIs. Easiest path is Manager:

Manager → Install via Git URL → https://github.com/kgonia/ComfyUI-KleinPrunedTE

or manually:

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

Then restart ComfyUI. Requires a recent ComfyUI with FLUX.2 support (Flux2Scheduler, EmptyFlux2LatentImage) - if you're on an older nightly, update first. You'll also need the models: the pruned encoder in models/text_encoders/, the gated flux-2-klein-9b.safetensors DiT in models/diffusion_models/ (accept BFL's license on Hugging Face and log in), and flux2-vae.safetensors in models/vae/. The example workflow's notes embed all three download links, so ComfyUI will offer to fetch the missing ones. Distilled Klein settings, already set in the example: 4 steps, CFG 1.0, euler.

Where people get burned

The big one is trying to load the pruned file with the stock CLIPLoader - don't. It can't express the pruned architecture and throws size mismatch errors. That's the point of this node, not a bug in it.

If you hit AttributeError: 'NoneType' object has no attribute 'device' from _apply_layout, that's the old-ComfyUI symptom: your build predates the FLUX.2/Klein code path the loader leans on. Update ComfyUI and restart.

Finally, mind the license: the pruned encoder is a derivative of FLUX.2-klein-9B and inherits the FLUX Non-Commercial License, same as the full model. The code here is MIT; the weights aren't.