--- language: - en license: other license_name: flux-non-commercial-license license_link: LICENSE.md base_model: black-forest-labs/FLUX.2-klein-9B tags: - flux - text-encoder - pruning - distillation - qwen3 - diffusers library_name: transformers --- # FLUX.2-klein-9B Text Encoder — Pruned 5.1B A structurally pruned drop-in replacement for the 8.2B Qwen3 text encoder of [FLUX.2-klein-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B): **8.19B → 5.10B parameters (−38%)**, recovered by hidden-state distillation against the original encoder. The DiT and VAE are untouched — this repo contains only the text encoder. | | original | **this repo** | this repo (fp8) | |---|---|---|---| | parameters (encode path) | 7.57B | **5.10B** | 5.10B | | weights | 14.1 GiB | **9.5 GiB** | ~4.8 GiB | | peak VRAM, **text encoder alone** (encode phase)¹ | 15.5 GiB | **10.6 GiB** | **6.8 GiB** | | embedding fidelity (masked token cos) | 1.0 | 0.9755 | 0.9750 | ¹ the encoder by itself, while encoding — whole-pipeline numbers below. **Whole-pipeline VRAM (this encoder fp8 + DiT fp8 + VAE), measured:** | configuration | resolution | peak VRAM | s/image | fits on | |---|---|---|---|---| | everything resident | 1024² | 16.4 GiB | 6.6 | 20 GB+ | | everything resident | 768² | 15.5 GiB | 4.2 | 16 GB (headless, tight) | | **DiT+VAE resident, encoder offloaded after encode** | 1024² | **11.0 GiB** | **6.5** | **12–16 GB** | | DiT+VAE resident, encoder offloaded | 768² | 10.1 GiB | 4.2 | 12 GB | Offloading only the encoder costs nothing: it runs once per prompt, and the DiT — the thing you don't want to swap — stays resident, so generation speed is unchanged (compare 6.5 vs 6.6 s/image). A fully CPU-offloaded pipeline (`enable_model_cpu_offload()`, bf16) manages ~31–39 s/image on the same GPU. With the **original bf16 encoder**, the fully-resident fp8-DiT setup needs ~26 GiB — this encoder is what brings it under the 24/20/16 GB thresholds. ## Overview FLUX.2-klein's pipeline consumes only three intermediate hidden states of its text encoder (layers 9/18/27 of 36) — the encoder is a feature extractor, not a language model. That structure makes large parts of it removable: 1. **Tail drop** — layers 28–35 are never read by the pipeline: removed exactly, no quality cost. 2. **Layer merge** — layers 10 and 19 SLERP-merged into their neighbors. 3. **FFN pruning** — activation-aware (Wanda) pruning of MLP width 12288 → 8192 on 7 layers. 4. **GQA head pruning** — whole key-value groups removed per layer (4–6 of 8 kept), guided by a per-layer sensitivity probe rather than a uniform budget. 5. **Export-time free removals** — the lm_head (622M, its output is never used for embeddings) and the final decoder layer whose output no tap reads. After each structural stage the model was **recovery-distilled against the original encoder** (never against a previous student — errors do not accumulate across stages) with per-tap hidden-state losses and a DiT-proxy loss through the frozen transformer. Export equality is enforced bitwise: the shipped model's prompt embeddings are `torch.equal` to the training-time checkpoint's. ## Results All 25 evaluation prompts, one fixed seed per row, generated side by side in a single session (4 steps, 1024×1024, guidance 1.0): Original | Pruned 5.1B :-------------------------:|:-------------------------: ![](images/original_00.png) | ![](images/pruned_00.png) ![](images/original_01.png) | ![](images/pruned_01.png) ![](images/original_02.png) | ![](images/pruned_02.png) ![](images/original_03.png) | ![](images/pruned_03.png) ![](images/original_04.png) | ![](images/pruned_04.png) ![](images/original_05.png) | ![](images/pruned_05.png) ![](images/original_06.png) | ![](images/pruned_06.png) ![](images/original_07.png) | ![](images/pruned_07.png) ![](images/original_08.png) | ![](images/pruned_08.png) ![](images/original_09.png) | ![](images/pruned_09.png) ![](images/original_10.png) | ![](images/pruned_10.png) ![](images/original_11.png) | ![](images/pruned_11.png) ![](images/original_12.png) | ![](images/pruned_12.png) ![](images/original_13.png) | ![](images/pruned_13.png) ![](images/original_14.png) | ![](images/pruned_14.png) ![](images/original_15.png) | ![](images/pruned_15.png) ![](images/original_16.png) | ![](images/pruned_16.png) ![](images/original_17.png) | ![](images/pruned_17.png) ![](images/original_18.png) | ![](images/pruned_18.png) ![](images/original_19.png) | ![](images/pruned_19.png) ![](images/original_20.png) | ![](images/pruned_20.png) ![](images/original_21.png) | ![](images/pruned_21.png) ![](images/original_22.png) | ![](images/pruned_22.png) ![](images/original_23.png) | ![](images/pruned_23.png) ![](images/original_24.png) | ![](images/pruned_24.png) ### Seed variance The 4-step distilled sampler is chaotic: the *same* prompt renders very differently across seeds — with either encoder — and any single seed can produce a degenerate draw. Two illustrations (top row = the seed used in the main table above): **"cyborg princess" — one seed of the four produced an artifact with the pruned encoder; the rest are clean for both:** Seed | Original | Pruned 5.1B :---:|:-------------------------:|:-------------------------: 4244 | ![](images/seed02_4244_original.png) | ![](images/seed02_4244_pruned.png) 11 | ![](images/seed02_11_original.png) | ![](images/seed02_11_pruned.png) 44 | ![](images/seed02_44_original.png) | ![](images/seed02_44_pruned.png) 66 | ![](images/seed02_66_original.png) | ![](images/seed02_66_pruned.png) **"extremely buff elon musk" — the prompt specifies no clothing; the encoders stably prefer different (both prompt-valid) interpretations:** Seed | Original | Pruned 5.1B :---:|:-------------------------:|:-------------------------: 4260 | ![](images/seed18_4260_original.png) | ![](images/seed18_4260_pruned.png) 11 | ![](images/seed18_11_original.png) | ![](images/seed18_11_pruned.png) 33 | ![](images/seed18_33_original.png) | ![](images/seed18_33_pruned.png) 55 | ![](images/seed18_55_original.png) | ![](images/seed18_55_pruned.png) **How to read the numbers.** Over 25 prompts, images conditioned by this encoder score SSIM ≈ 0.59 against images conditioned by the original — visibly *different renders* of the same prompt, not degraded ones. For calibration: the **original encoder against itself on two different GPUs scores SSIM 0.36** on identical prompts and seeds. This encoder diverges from the original *less than the original diverges from itself across hardware*. | comparison (25 prompts, same session) | SSIM | MSE | LPIPS | |---|---|---|---| | pruned (bf16) vs original | 0.591 | 0.030 | 0.306 | | pruned (fp8) vs original | 0.591 | 0.030 | 0.311 | | original vs original, different GPU | 0.357 | 0.058 | 0.529 | ## Benchmarks Measured on an RTX 5090 (24 GB), torch 2.11 / cu128, 512-token encodes, batch 4: | variant | params | weights | encode peak VRAM | masked token cos vs original | |---|---|---|---|---| | original bf16 | 7.57B | 14.1 GiB | 15.47 GiB | 1.0 | | original fp8 | 7.57B | ~7.1 GiB | 9.38 GiB | 0.9987 | | **pruned bf16** | 5.10B | 9.5 GiB | 10.56 GiB | 0.9755 | | **pruned fp8** | 5.10B | ~4.8 GiB | **6.77 GiB** | 0.9750 | ## Quick Start > **Prerequisite:** the DiT/VAE/tokenizer come from the gated base repo — accept the > license at [black-forest-labs/FLUX.2-klein-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) > and `hf auth login` once. This repo only replaces the text encoder, so you never > download the original 14 GiB encoder — total download is the same as the base > pipeline alone (~27 GiB), and less if you already have klein cached. The encoder has mixed per-layer FFN/attention widths, which a stock Qwen3 config cannot express — load it through the bundled `loading.py` (plain file, no `trust_remote_code`): ```python import torch from huggingface_hub import snapshot_download repo = snapshot_download("SearchingMan/FLUX.2-klein-9B-Text-Encoder-Pruned-5.1B") import sys; sys.path.insert(0, repo) from loading import load_text_encoder, load_pipeline te = load_text_encoder(repo, torch_dtype=torch.bfloat16) pipe = load_pipeline("black-forest-labs/FLUX.2-klein-9B", te, torch_dtype=torch.bfloat16) pipe.enable_model_cpu_offload() image = pipe( "A cat holding a sign that says hello world", text_encoder_out_layers=(9, 17, 25), # REQUIRED — see warning below num_inference_steps=4, guidance_scale=1.0, height=1024, width=1024, generator=torch.Generator("cuda").manual_seed(0), ).images[0] ``` > ⚠️ **Every pipeline call must pass `text_encoder_out_layers=(9, 17, 25)`.** > The pipeline's default taps assume the original 36-layer encoder; without this argument > it will silently read the wrong hidden states and produce degraded images. > Never load this checkpoint with `ignore_mismatched_sizes=True`. **fp8 (recommended on consumer GPUs):** ```python from torchao.quantization import quantize_, Float8WeightOnlyConfig quantize_(te, Float8WeightOnlyConfig()) # encoder → ~4.8 GiB quantize_(pipe.transformer, Float8WeightOnlyConfig()) # DiT → ~9 GiB pipe.to("cuda") # fully resident, ~16.4 GiB peak ``` ## Limitations - **Non-commercial license** (inherited from FLUX.2-klein-9B — see LICENSE.md), and you need access to the gated base repo for the DiT/VAE/tokenizer. - This is a **component**, not a standalone model: it produces prompt embeddings for FLUX.2-klein-9B only. - Images conditioned by this encoder are **different renders**, not pixel-matched ones: composition details (colors, props) can flip relative to the original on a given seed — the same class of change you get from running the original on different hardware. In a seed study on artifact-suspect prompts, anatomical errors occurred at the same seed-dependent rate as with the original encoder (0 in 63 fresh student draws vs 0 in 21 original draws) — if a render shows one, regenerate with a new seed. On some prompts the pruned encoder consistently prefers a different, still prompt-faithful interpretation (see the seed-variance examples above). - Same-seed outputs are only comparable within one GPU/software environment (a property of the 4-step distilled sampler, not of this encoder). - The encode protocol is fixed: Qwen3 chat template, `enable_thinking=False`, 512-token max length — handled automatically by the pipeline. ## Training details | stage | what | recovery | |---|---|---| | tail drop | layers 28–35 removed | exact, none needed | | SLERP merge | 10→11, 19→20 (α=0.7) | 3500 steps | | FFN Wanda pruning | 7 layers, 12288→8192 | 3500 steps (combined with merge recovery) | | GQA head pruning | 13 layers, keep 4–6 of 8 groups, probe-guided | 2000 steps | | export | delete merged/tail/final layers, drop lm_head, remap taps (9,18,27)→(9,17,25) | bitwise-verified | Distillation: ~40k captions (a general text-to-image prompt corpus plus text-rendering-focused prompts), losses on masked per-tap hidden states (cosine + normalized MSE + norm + Gram), plus a DiT-proxy loss through the frozen transformer at high-noise timesteps. Teacher was always the original encoder. Head-pruning budgets came from a 42-point per-layer sensitivity probe; the pruned groups are pinned in `pruning_metadata.json`. ## Provenance & license This is a **modified version (Derivative) of the FLUX.2-klein-9B text encoder** by Black Forest Labs, distributed under the FLUX Non-Commercial License (see `LICENSE.md` and `NOTICE`). This project is not affiliated with, endorsed, approved, or validated by Black Forest Labs. For commercial licensing of FLUX models see https://bfl.ai/licensing. ## Support Models like this one are trained on my own hardware and my own cloud budget. If this work saves you VRAM, time or money, you can [buy me a coffee](https://ko-fi.com/michelangelofussion) or simply follow the next runs on [X](https://x.com/kgonia7) and like this model. Every bit funds the next experiment. Full write up: [kgonia.github.io](https://kgonia.github.io/projects/flux2-klein-text-encoder-pruned/)