Skip to content

feat(flux3): add FLUX.3 Image Klein and Klein Base - #9883

Open
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/flux3-klein
Open

Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/flux3-klein

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

Important

Do not merge before a diffusers release ships Flux3Transformer2DModel. Review is welcome now.

Flux3Transformer2DModel exists only on diffusers main (commit 7ecb1e1b4f, "Flux3 image (#15046)"). The PR vendors that file, unchanged apart from absolute imports, so it runs and can be tested on the pinned diffusers==0.41.0. Before merging:

  1. bump the diffusers pin to the release;
  2. delete invokeai/backend/flux3/transformer_flux3.py;
  3. import the class from diffusers in invokeai/backend/flux3/__init__.py, the only import point;
  4. drop its ruff format exclusion in pyproject.toml.

tests/backend/flux3/test_vendored_transformer.py fails on the bump and lists these steps.

Summary

Adds Black Forest Labs' FLUX.3 Image Klein as a new flux3 base with two variants:

  • Klein: step-distilled flow map, 16 steps, no CFG.
  • Klein Base: 50 steps, CFG 3.5, negative prompt.

Both handle text-to-image and editing with up to 10 reference images. They install as Diffusers, as a bf16 single file or as GGUF.

Backend

  • flux3 base and Flux3VariantType.
  • Identification of the Diffusers folder (modular_model_index.json), the bf16 single file and GGUF, with FLUX.1/FLUX.2 exclusions; community GGUFs declare general.architecture="flux".
  • Loaders, native-key conversion, Qwen3-VL 4B text encoding, sampling, and reference encoding.
  • Nodes flux3_model_loader, flux3_text_encoder and flux3_denoise (Prototype).
  • Starter models and a bundle (ungated: GGUF Q4_K_M, FLUX.2 VAE, Qwen3-VL 4B).
  • FLUX.3 latents are FLUX.2 latents (byte-identical VAE), so the FLUX.2 VAE encode/decode and PiD nodes are reused.
  • Quantized single files (fp8, int8, nvfp4) are refused at install with a reason until a follow-up adds them.
  • Installing by repo id now reads modular_model_index.json, so it downloads only the Diffusers folders, not the 11 GB root single file next to them.

webv2

  • Generate, reference editing, and canvas img2img/inpaint/outpaint.
  • Reference cards show the <ref_image_N> tag the prompt uses, numbered like the graph's reference order; clicking the tag copies it.
  • Klein with CFG > 1 blocks Invoke with a reason.
  • FLUX.2 PiD decoders work on FLUX.3.

Docs

  • New FLUX.3 page.
  • FLUX.3 added to the model-families, format-support, system-requirements and Generate pages.

Follow-ups (stacked, not in this PR): layout boxes from canvas regions, quantized single files, LoRA, and BFL's prompt upsampler.

Related Issues / Discussions

None. First PR of the FLUX.3 stack.

QA Instructions

Automated (Windows, RTX 4090):

  • uv tool run ruff@0.11.2 check / format --check on all changed Python: clean.
  • uv run --no-sync pytest -n 4 over tests/backend/flux3, tests/backend/flux2, tests/backend/architectures, tests/backend/model_manager/configs, tests/app/invocations and tests/backend/util and tests/app/routers/test_architecture_capabilities.py: 4114 passed.
  • tests/backend/util/test_sage_attention.py and tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py fail to collect on main too ('filterwarnings' not found in markers); they are unrelated to this PR.
  • webv2:
    • format:check, lint:oxc, lint:tsc and architecture:check pass.
    • pnpm test: 10202 passed, 3 failed. The three fail only under a German locale (digit grouping in clusterStats, indexProgress and modelCacheStore) and are unrelated.
  • invokeai/frontend/api: lint and tests pass.
  • docs: node --test passes and pnpm build succeeds.

Numerical comparison against the diffusers main reference pipeline with BFL's noise:

  • Setup: 1024², bf16, use_cache=false, model cache emptied between models.
  • The InvokeAI conversion is bit-identical to diffusers' own (359/359 tensors).
Case PSNR vs diffusers
Klein t2i, seeds 1/2/3 49.3 / 48.4 / 49.3 dB
Klein, 1 reference, seeds 1/2/3 51.0 / 50.3 / 46.3 dB
Klein, 3 references (768² each) 42.0 dB
Klein Base t2i, CFG 3.5, seeds 1/2/3 36.0 / 36.4 / 36.8 dB
bf16 single file + fp8 Qwen3-VL encoder 43.7 / 38.4 dB
GGUF Q8_0 42.0 / 39.7 dB
GGUF Q4_K_M 24.1 / 27.5 dB

Timings, RTX 4090, 1024²:

  • Klein: about 15 s.
  • Klein with 1 reference: about 24 s.
  • Klein with 3 references: about 31 s (768² references).
  • Klein Base: about 51 s.

Memory, measured at driver level on the 4090:

  • Transformer forward:
    • 0.286 MB per attended token without references;
    • 0.320 MB per token with six 1 MP references;
    • the node budgets 0.33 MB/token plus 1 GB.
  • Untiled reference VAE encode: 2.07 GB at 1024² and 8.26 GB at 2048², against estimates of 2.15 and 8.59 GB.

Browser (webv2), all checked:

  • Klein t2i with two references, prompted with their tags.
  • Disabling a reference card renumbers the tags of the cards after it.
  • Tag copy: success feedback, and the error toast when no clipboard is available.
  • Klein at CFG 2 blocks Invoke with the reason.
  • Base defaults (50 steps, CFG 3.5) arrive through "Reset all to model defaults".
  • Canvas img2img produces the expected graph (flux2_vae_encode → flux3_denoise, FLUX.2 strength curve).
  • FLUX.2 PiD 1.5 decode of FLUX.3 latents (Fit to size).

To try it, install the FLUX.3 bundle (or black-forest-labs/flux-3-image-klein, which is gated) and pick FLUX.3 Klein on the Generate tab.

Review

Material findings resolved:

  • Reference encoding. Tiling the reference VAE encode moved every reference token, not only the seams; 1-reference edits matched diffusers at only 30–34 dB. FLUX.3 now encodes references untiled when fused SDPA keeps the encode linear and the estimate fits half the device memory. They are still tiled on ROCm and MPS, and for large references on small cards.
  • Tokenizer sizing. In a full repo download, the tokenizer was sized as the 11 GB root single file and could evict the transformer.
  • GGUF. Unquantized tensors are now unpacked at load through unpack_ggml_at_load instead of being decoded every forward. QK norms keep their stored fp32 values.
  • Klein with CFG is now refused before any model loads.
  • Node wiring tests. They cover the img2img normalization rule ([bug]: Flux.2 Klein Img2Img Posterization #8964), the inpaint noise and mask inversion, schedule respacing, and the degenerate-window return.
  • webv2 tests. They cover the reference order versus the prompt tags, the negative branch per variant, and the Klein CFG validation.
  • Smaller fixes:
    • the reference size limits now hold at the edges;
    • the compute dtype follows the device;
    • the head of a split denoise chain is not respaced;
    • the tag chip gives copy feedback.

Remaining risks and limitations:

  • Working memory per token is measured on CUDA only; ROCm is unmeasured.
  • The 12 GB minimum VRAM in the system requirements is an estimate.
  • FLUX.2 still tiles its reference encode and has the same quality loss. Changing that is a FLUX.2 behavior change, left for its own PR.
  • SageAttention is not wired into flux3_denoise.
  • PiD decode is not available on the canvas; this is an existing webv2 limitation.

Compatibility / Rollout

  • Persisted data. The new flux3 base and the flux3_klein / flux3_klein_base variants are additive; existing model records are unaffected and no migration is needed.
  • API. The OpenAPI schema and the generated schema.ts change by additions only: three nodes, one conditioning field and output type, enum values, and the flux3_* generation modes.
  • Shared FLUX.2 code. Changes are behavior-preserving refactors:
    • BN statistics and normalization moved to backend/flux2/latent_norm.py;
    • reference encoding moved to encode_reference_image, whose default keeps FLUX.2's tiled encode;
    • Krea-2's Qwen3-VL variant check is shared.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

@github-actions github-actions Bot added python PRs that change python files Root invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend PRs that change frontend files python-tests PRs that change python tests docs PRs that change docs python-deps PRs that change python dependencies labels Oct 16, 2026
@Pfannkuchensack Pfannkuchensack mentioned this pull request Oct 16, 2026
1 task done
submodels = resp.json()
paths = [Path(subfolder or "", x) for x in paths if Path(x).parent.as_posix() in submodels]
paths.insert(0, Path(f"{prefix}model_index.json"))
elif not prefix and Path("modular_model_index.json") in paths:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Account for the repository prefix when filtering modular downloads

HuggingFaceMetadataFetch.from_id() constructs paths as Path(name, s.rfilename), and filter_files() preserves that prefix. Consequently, an actual install of black-forest-labs/flux-3-image-klein contains flux-3-image-klein/modular_model_index.json, so this condition is always false. The install retains the approximately 11 GB root checkpoint and weights under variants/ alongside the Diffusers components. A reproduction using the pinned selector and prefixed paths confirms both unwanted checkpoints remain selected. Normalize paths relative to the repository root for the index check and component-folder comparison, while retaining the original paths in the returned files. The added test currently uses unprefixed paths, so it does not exercise the production metadata contract.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks. from_id() builds paths as Path(name, rfilename), so the modular branch never matched. Fixed in 045e7e3: the modular_model_index.json check and the component-folder comparison now work relative to the repository root, while the returned files keep their prefixed paths. The test now uses production-shaped prefixed paths and fails against the previous code.

The classic model_index.json branch has the same pre-existing prefix issue (it never runs for from_id metadata either); that is out of scope here and will get its own issue.

Adds the flux3 base (Klein flow map, Klein Base) with a vendored transformer, Diffusers/bf16/GGUF loading, text encoder, denoise node with up to 10 reference images, starter models and docs.
webv2 generates, edits and runs canvas img2img/inpaint/outpaint with FLUX.3, shows <ref_image_N> tags on reference cards and blocks CFG on Klein.
FLUX.3 encodes references untiled where they fit; quantized single files are refused at install until a follow-up.
…o root

Metadata paths carry the repository name, so the modular_model_index.json check never matched and a repo-id install also pulled the root single file and variants/.
# File paths start with the repo name (HuggingFaceMetadataFetch.from_id), so folders compare below it.
url = hf_hub_url(self.id, filename="modular_model_index.json")
try:
resp = session.get(url, timeout=HF_METADATA_REQUEST_TIMEOUT_SECONDS)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Authenticate the modular model-index request

Installing either gated Black Forest Labs FLUX.3 Diffusers repository reaches this request without an Authorization header, including through the new starter models. ModelInstallServiceDefault._import_from_hf() puts the configured/cached token on source.access_token, but _remote_files_from_source() does not pass that token to download_urls(), which normally creates a plain requests.Session; the token is only used later by _multifile_download(). Consequently, fetching modular_model_index.json fails with an authorization error before the download job is created, even after the user accepts the license and configures a valid token. Use an authenticated Hub request or propagate the source token into this metadata fetch, with coverage for the gated installation path.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend PRs that change backend files DO NOT MERGE docs PRs that change docs frontend PRs that change frontend files invocations PRs that change invocations python PRs that change python files python-deps PRs that change python dependencies python-tests PRs that change python tests Root services PRs that change app services

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants