FLUX.2-klein-9B-Text-Encoder-Pruned-5.1B

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: 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.

¹ the encoder by itself, while encoding — whole-pipeline numbers below.

Whole-pipeline VRAM (this encoder fp8 + DiT fp8 + VAE), measured:

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

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:

"extremely buff elon musk" — the prompt specifies no clothing; the encoders stably prefer different (both prompt-valid) interpretations:

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.

Benchmarks

Measured on an RTX 5090 (24 GB), torch 2.11 / cu128, 512-token encodes, batch 4:

Quick Start

Prerequisite: the DiT/VAE/tokenizer come from the gated base repo — accept the license at 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):

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

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

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 or simply follow the next runs on X and like this model. Every bit funds the next experiment. Full write up: kgonia.github.io