README.md

---
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/)