Welcome FLUX.3 Image Klein in Diffusers

What it claims

FLUX.3 Image Klein ships in two open checkpoints: black-forest-labs/flux-3-image-klein (distilled) and black-forest-labs/flux-3-image-klein-base. A single model handles text-to-image, multi-reference editing and bounding-box layout. The distilled checkpoint defaults to 16 steps at 1024×1024. The base checkpoint runs 50 steps with classifier-free guidance and is the one intended for fine-tuning.

How it works

The pipeline uses the FLUX.2 VAE, which downsamples 8× into 32 channels. Each 2×2 latent patch is then packed into one 128-channel token, so one token covers 16×16 pixels and a 1024² image becomes 4,096 tokens. Reference images go through the same encoder and keep their own size.

Text is encoded by Qwen3-VL. The model takes hidden states from 8 of its 36 layers (layers 4, 8, …, 32), stacks them into a 20,480-dim vector per token, and projects that to a width of 3072. Prompts are padded to a multiple of 80 tokens, and the transformer attends to that padding. This is why the pipeline accepts only one prompt per call.

flowchart LR
  T[Text tokens] --> ST[5 text-only blocks]
  N[Noisy image tokens] --> SN[5 image-only blocks]
  R[Reference tokens] --> SR[5 reference-only blocks]
  ST --> J[28 shared joint blocks<br/>full attention]
  SN --> J
  SR --> J
  J -->|image tokens only| O[Output layers → velocity]
  J -.->|text & refs dropped| X[discarded]

The transformer runs in two stages:

• Stage 1: each stream gets 5 blocks of its own weights and attends only to itself.
• Stage 2: the streams are concatenated and pass through 28 shared blocks with full attention. FLUX.2, by contrast, let text and image attend to each other from the start.

Each block runs attention and a SwiGLU MLP in parallel on the same normalized input. Scale, shift and gate come from timestep modulation, set separately for text, image and reference tokens. Image tokens get the current timestep t, while text and reference tokens always get t=0, which marks them as clean context. The distilled checkpoint is a flow map: it also embeds the target timestep of each jump.

RoPE uses four axes (t, h, w, l), each taking 32 of the 128 head dims. References share (h, w) coordinates with the output, which is what makes src_bbox/tgt_bbox editing precise. The scheduler applies a time shift of 6.93, so 44 of 50 steps run above σ=0.5.

Evidence

• The distilled checkpoint makes a 1024² image in about 2 s on a single H100 80GB. The setup is 16 steps, no offloading, default attention, averaged over 3 runs after warmup.
• Adding one 1024² reference roughly doubles generation time, because it adds 4,096 tokens to every transformer call.
• CPU offload was not benchmarked.

How to use it

pip install -U diffusers transformers

pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein")
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")
pipe(prompt="a red fox", output="images")[0].save("fox.png")

Prompt conventions:

• Multiple references: pass image=[a, b] and refer to them in the prompt as <ref_image_0>, <ref_image_1>.
• Layout: append a JSON list of {id, bbox, desc} after the caption, and mention each region in the caption as <id>.
• Box format: boxes are [y_min, x_min, y_max, x_max], y first, on a 0–1000 scale that is independent of resolution.
• Editing: each region uses from, src_bbox and tgt_bbox. Setting both boxes equal pins that region in place.

Base checkpoint: load with dtype={"vae": torch.float32, "default": torch.bfloat16}. Guidance is optional, e.g. pipe.guider.new(guidance_scale=5.0).

Quantized transformers come in fp8r (rowwise FP8) and nvfp4. Load them with Flux3Transformer2DModel.from_single_file(...) and BFLQuantizationConfig, then call pipe.update_components(transformer=...). This avoids also downloading the bf16 transformer.

To save memory, pop the text_encoder block, encode the prompt once, free the encoder, then generate from prompt_embeds. Alternatively, use ComponentsManager().enable_auto_cpu_offload(...).

LoRA training runs on the base checkpoint:

accelerate launch train_dreambooth_lora_flux3_klein.py \
  --pretrained_model_name_or_path="black-forest-labs/flux-3-image-klein-base" \
  --dataset_name="multimodalart/1920-raider-waite-tarot-public-domain" \
  --caption_column="caption" --instance_prompt="trtcrd" \
  --mixed_precision="bf16" --resolution=1024 --use_aspect_ratio_buckets \
  --train_batch_size=1 --rank=16 --lora_alpha=16 --learning_rate=1e-4 \
  --lr_scheduler="constant" --max_train_steps=1500

• To cut VRAM, add --gradient_checkpointing --cache_latents --use_8bit_adam.
• For editing, train_dreambooth_lora_flux3_klein_img2img.py trains on (reference, target, instruction) triplets.
• Ostris' AI Toolkit is also supported.

Limits

• One prompt per call, because the padding is attended to.
• Each reference image adds latency that scales with its token count.
• The page gives no licence, parameter count, quality benchmarks or memory figures, so check the model cards before adopting.
• Offloading performance is unmeasured.