Skip to content

FLUX 3 Image Klein: text-to-image and single-reference editing - #73

Merged
TroyHernandez merged 4 commits into
mainfrom
feat/flux3-klein
Oct 28, 2026
Merged

TroyHernandez merged 4 commits into
mainfrom
feat/flux3-klein

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

Adds FLUX 3 Image Klein (black-forest-labs/flux-3-image-klein): text-to-image and single-reference editing, ported from the diffusers pipeline. Version 0.2.2.16.

download_flux3_klein()   # gated: accept the FLUX Kommunity License + HF_TOKEN; ~20 GB, one-time fp8 quantize to 5.3 GB
pipe <- flux3_load_pipeline()
txt2img_flux3(caption, pipeline = pipe, seed = 7)
img2img_flux3(caption, image = "photo.png", pipeline = pipe, seed = 7)

What's in it

  • flux3_transformer(): per-mode blocks (image, reference, text; 5 each) then 28 joint blocks, per-token timestep modulation, target-timestep embedder. Built as a skeleton so a bf16 load peaks at ~15 GB host RAM instead of 22 GB.
  • Qwen3-VL-4B text model through load_qwen3_text_encoder() (layers 4..32 stacked, prompt padded to a multiple of 80).
  • flux3_preprocess_image(): reference kept whole, scaled to fit max_area (default 4 MP), reflect-padded to the 16 px grid, output cropped back.
  • fp8 quantization (cast census 265) and the gated download, pinned to revision 3373376.
  • diffuseR takes the caption as given. BFL recommends a prompt upsampler, and expanding short prompts is the caller's job.

Shared fixes

  • VAE mid-block attention now uses fused SDPA. A 2048x2048 image allocated an 8 GiB score matrix and ran out of memory on 16 GB. This code is shared by the SD, SDXL, FLUX.1, FLUX.2 and FLUX 3 VAEs.
  • The Qwen3 encoder is built at its target dtype, which removes a 16 GB fp32 intermediate when loading in bf16. Requesting the last layer now returns the post-norm state, matching transformers.
  • flux2_feed_forward uses linear_noinit.

Validation (real weights)

  • Tokens and prompt template are identical to the shipped processor. Qwen3-VL embeddings are bit-identical at fp32.
  • Transformer at fp32, depth-reduced but covering every block type: max abs error 2.9e-6, with and without a reference.
  • 1 MP edit on the GPU against diffusers, each with its own text encoder: R bf16 52.2 dB PSNR, R fp8 36.7 dB.
  • RTX 5060 Ti 16 GB, fp8:
    • 1024² text-to-image: 44 s, peak 11.7 GB allocated.
    • 1 MP edit: 50 s, peak 13.8 GB allocated.
    • 4 MP edit: 285 s, peak 14.1 GB allocated.
  • tinytest: 1540 passed, 0 failed (53 new).

Not in this PR

  • An 8 GB mode (fp8 text encoder plus a block-streamed transformer).
  • Multiple references.
  • Fleet entries.

The weights are under the FLUX Kommunity License and are never redistributed; see inst/COPYRIGHTS.

Port of the diffusers FLUX 3 transformer and distilled modular pipeline:
flux3_transformer (per-mode + joint blocks, per-token modulation,
flow-map target timestep), Qwen3-VL text path (flux3_encode_prompt; the
Qwen3 encoder gains position ids, Qwen3-VL loading, and HF's post-norm
last hidden state), fill-mode reference preprocessing, txt2img_flux3,
img2img_flux3, flux3_load_pipeline and download_flux3_klein (gated,
pinned snapshot, fp8 census 265). FLUX 3 transformers load as skeletons
to keep host RAM under the bf16 size.
@TroyHernandez
TroyHernandez merged commit f5cab69 into main Oct 28, 2026
2 checks passed
@TroyHernandez
TroyHernandez deleted the feat/flux3-klein branch October 28, 2026 19:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant