Feat: flux2 dev support

[Pfannkuchensack]

Summary

Adds end-to-end support for FLUX.2 [dev] alongside the existing FLUX.2 Klein implementation. Dev uses a Mistral Small 3 text encoder (hidden 5120 → joint_attention_dim 15360, hidden states sampled at layers 10/20/30) instead of Klein's Qwen3. It shares the 32-channel AutoencoderKLFlux2 VAE and the 4D-RoPE sampling backend with Klein, so most existing infrastructure is reused — only the Mistral encoder loaders/configs and the Dev graph wiring are new.

Two Mistral encoder packagings are auto-detected: the 40-layer Mistral that ships in the black-forest-labs/FLUX.2-dev diffusers repo, and the 30-layer cow-mistral3-small distillation (Comfy-Org safetensors, gguf-org GGUFs). Single-file and GGUF encoders embed the Tekken tokenizer, so no separate tokenizer download is needed.

What works

• Generation modes: txt2img, img2img, inpaint, outpaint (canvas)
• Multi-reference image editing via the FLUX Kontext ref-image UI (same as Klein)
• LoRAs for FLUX.2 [dev]
• Model formats: diffusers folder, single-file safetensors, and GGUF transformers; Mistral encoder as diffusers / single-file (bf16 / fp8 / fp4) / GGUF, with standalone VAE + encoder selection in Advanced Settings
• Metadata recall / Remix: Dev VAE + Mistral encoder restore correctly, disambiguated from Klein's VAE / Qwen3 slice
• Starter models for the Dev transformer, VAE, and Mistral encoders

Limitations / notes

• Memory: the Mistral encoder is large (fp8 ~18 GB, bf16 ~35 GB) — plan VRAM/RAM accordingly.
• Encoder precision matters: very low-bit GGUF Mistral encoders (e.g. Q2/Q3) noticeably weaken prompt adherence; prefer fp8 / bf16 or Q6/Q8.
• Non-BFL Mistral: upstream Mistral Small 3.1 / 3.2 (40-layer, non-BFL) is accepted but off-distribution for FLUX.2 and logs a load-time warning; it is not offered as a starter model.
• Tokenizer fallback: encoders that don't embed Tekken (older cow GGUFs, bare diffusers folders without tokenizer/) fall back to fetching black-forest-labs/FLUX.2-dev:tokenizer from HF; offline with no HF cache this errors with documented workarounds.
• License: FLUX.2 [dev] is Non-Commercial (flagged in isNonCommercialMainModelConfig).
• Applying a Klein LoRA on a Dev model logs a variant-mismatch warning (no crash).

How to test

Automated:

# backend (fast, no model load)
uv run --extra cuda pytest tests/test_imports.py tests/model_identification/test_identification.py

pyproject.toml gained mistral-common; uv.lock is regenerated (uv lock --locked passes).

End-to-end:

1. In the Model Manager install a Dev transformer (a gguf-org/flux2-dev-gguf quant — ~20 GB for Q4_K_M — or the full black-forest-labs/FLUX.2-dev diffusers folder), the FLUX.2 VAE, and a Mistral encoder (Comfy-Org fp8 recommended, or a Q6/Q8 cow GGUF).
2. Select the Dev transformer. Advanced Settings should show FLUX.2 [dev] VAE and FLUX.2 [dev] Mistral Encoder dropdowns (not Klein's Qwen3 selector). A missing VAE/encoder should block queueing with a Dev-specific message.
3. Defaults: 28 steps, guidance 3.5, CFG 1.0, 1024×1024. Generate — the first run with Comfy / cow files logs Loaded embedded Tekken tokenizer from <file> (no HF fetch).
4. Exercise img2img, inpaint / outpaint on canvas, and a multi-reference edit.
5. Apply a FLUX.2 [dev] LoRA and generate.
6. Recall: open a generated Dev image → Remix; Dev VAE + Mistral encoder should restore. Repeat with a Klein image and confirm Klein VAE + Qwen3 still recall (no cross-contamination).

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

Closes #8668

[lstein]

Frontend build is failing. pytests are failing.

[lstein]

I've re-verified the branch at 3e62aa8 against my second review. Thank you for the thorough fix commit (5026c02) — 8 of the 10 findings are correctly and completely fixed, including all the hard parts (the migration seed with a field-accurate test fixture, model-resolving guidance recall, the _reinit_inv_freq extraction validated against transformers 5.5.4, and the vision-tower drop). Two of the fixes need a follow-up, which is why this is another change request rather than an approval. Details below.

Requested changes

🔴 1. The new Klein-loader variant guard broke the VAE-only fallback (regression from the blocker-3 fix)

flux2_klein_model_loader.py calls _validate_diffusers_format on both call sites — the encoder-extraction path (:154) and the VAE-only extraction path (:132). The encoder-side rejection is exactly right, but the 32-channel AutoencoderKLFlux2 is shared between Klein and [dev], and the linear UI deliberately relies on that: buildFLUXGraph.ts falls back to any FLUX.2 diffusers pipeline when only the VAE is needed ("Fall back to any diffusers model if only the VAE is needed"), and isFlux2DiffusersMainModelConfig / readiness don't filter by variant.

Triggering sequence: install a Klein GGUF main + a standalone Qwen3 encoder + a FLUX.2 [dev] diffusers pipeline as the only FLUX.2 diffusers model (no standalone VAE) → readiness enables Invoke → the builder sets qwen3_source_model to the [dev] pipeline for VAE extraction → the new guard raises ValueError claiming the VAE is "incompatible with the Klein transformer" — which worked before this fix, and the message is factually wrong for the VAE.

Suggested fix: apply the variant guard only at the tokenizer/encoder extraction call sites ([dev] loader :153, Klein loader :154), not the VAE ones — or, if you prefer to reject cross-variant VAE sourcing on principle, the frontend readiness + source-selection fallback must be variant-filtered to match, so users get a disabled Invoke button with a reason instead of a runtime error.

🔴 2. The root-directory tokenizer step still fails offline for the exact layout blocker 4 described

The structural cleanup is all there (dead fallback gone, duplicated pre-try gone, ValueError in every except tuple, all three loaders share the helper). But the new step 3 (mistral_encoder.py:676-681) tries AutoProcessor only. With the pinned transformers 5.5.4, AutoProcessor.from_pretrained on a directory whose config.json is model_type: "mistral3" — which is precisely the BFL-style standalone-encoder layout this step targets — raises OSError: Can't load image processor … (it resolves mistral3 to a multimodal processor and requires preprocessor_config.json) before reaching any tokenizer fallback. The exception is swallowed and the ladder falls through to the HF fetch, which raises RuntimeError offline.

Reproduced empirically against the project venv with HF_HUB_OFFLINE=1: a root directory with a mistral3 config.json plus a valid tekken.json (or tokenizer.json + tokenizer_config.json) fails via AutoProcessor, while AutoTokenizer.from_pretrained on the same directory succeeds and returns a working backend. So the tokenizer is loadable — only the loader-class choice loses it. The step only succeeds today when the root config is text-only model_type: "mistral" or a full processor set including preprocessor_config.json is present.

Suggested fix: loop (AutoProcessor, AutoTokenizer) in steps 2 and 3, exactly as _load_tokenizer_from_hf already does (:616). Note if you do: probing showed AutoTokenizer can raise KeyError: 'special_tokens' on a directory containing only tekken.json, so the except tuple there needs KeyError too (or that escapes loudly).

Note on the two commits I just pushed (merge of main + migration v5 bump)

I've pushed 567cba9 (merge of latest main, which brings in the PiD decoder work — the flux2 PiD wiring is adapted into this branch's addFlux2Features helper and verified behavior-identical to main's) and d3c7bde. The second commit fixes a semantic conflict the textual merge couldn't see, and it wasn't your bug: main (after the PiD PR, 802de41) and this branch both bumped the params persist schema to _version: 4 with different new keys (main: pidMode/pidDecoderModel/gemma2EncoderModel/pidSteps, still with kleinVaeModel; this branch: flux2VaeModel + flux2DevMistralEncoderModel). After a plain merge, a v4 blob written by either parent is missing the other side's keys, skips every migration branch, fails zParamsState.parse(), and silently wipes the whole params slice — the same failure mode as blocker 1, reintroduced by the version collision.

The fix keeps main's v3→v4 verbatim and moves this branch's flux2 slot-merge + Mistral seed into a new v4→v5 step with conditional (?? null) seeding for both v4 shapes, with tests covering both. It also seeds the five wan* keys in v3→v4: main added them to the schema without a version bump while v6.13.x releases were still writing v3 blobs, so genuine released-build v3 blobs fail parse on those keys today (empirically confirmed — inherited from main, needs the same fix there).

Verified fixed (no action needed)

1. Migration seed (blocker 1) — flux2DevMistralEncoderModel = null seeded in the migration; the rebuilt fixture genuinely covers the regression (deleting the seed fails the test). ✅
2. Guidance for [dev] (blocker 2) — the parse gate resolves the image's own model and exempts variant === 'dev'; Klein still skips; slider renders for dev only via the same predicate the graph builder uses, so visibility and consumption can't drift; generate → metadata → recall → regenerate round-trips. Uninstalled-model dev images fall back to skipping guidance — reasonable fail-safe. ✅
3. Encoder-side variant guards (blocker 3) — the reported Klein-encoder-into-dev-transformer path (and mirror) is closed on every reachable route with clear ValueErrors, and the linear UI filters dev sources by variant. (VAE-path over-reach above is the only issue.) ✅
4. Tokenizer structure (blocker 4) — dead fallback and duplicate pre-try removed, ValueError caught, single helper used by all three loaders. (AutoProcessor-only gap above is the only issue.) ✅
5. Collection LoRA base check — ValueError before dedup/config fetch; FLUX.1 rejected; round-1 Klein-variant rejection intact; Klein side symmetric. ✅
6. Mistral slot clearing — cleared on base switch (with toast accounting) and intra-flux2 variant switch, including when the Klein encoder is unset; recall path covered via modelSelected. ✅
7. rope_theta — single _reinit_inv_freq with the rope_parameters/rope_scaling fallback; verified on transformers 5.5.4 that rope_theta set via MistralConfig(rope_theta=…) lands exactly where the fallback reads it. ✅
8. mistral-common pin — >=1.5.4,<2 in the right block, lock resolves 1.11.6, and both private-API usages (_special_tokens_reverse_vocab, Tekkenizer.encode(s, bos, eos)) verified present in-range with defensive fallbacks. ✅
9. Vision tower drop — setattr(model, …, None) removes the tower/projector from state_dict()/named_parameters(), so cache accounting and VRAM transfers exclude it; .language_model access and index-based extraction unaffected. Residual, non-blocking: the sibling full-pipeline path (flux.py Flux2DiffusersModel._load_model, used for diffusers mains and mistral_source_model) still loads and carries the ~0.8 GB tower — worth the same treatment in a follow-up. ✅
10. 40-layer docs — all four sources now agree on the operational contract (accepted, warned, indices 10/20/30, final norm kept for 40-layer / stripped for cow). Nit, non-blocking: starter_models.py says BFL's own 40-layer has weaker adherence due to index depth, while the loader warning implies only non-BFL 40-layer weights degrade — pick one story. ✅

Pre-existing, out of scope for this PR (noting for completeness): the workflow editor can still feed a Klein 9B-family source into a 4B main (no cross-Klein variant check on the source-extraction path — same wrong-width class, linear UI is protected); the Klein loader never validates its main model's variant.

🤖 Generated with Claude Code

[lstein]

Round 4 review (head c2b7778)

Thanks for the quick turnaround on both round-3 blockers. I re-verified each fix in code and with fresh-context adversarial passes (including mutation-testing your new tests — they do catch reverts). Summary:

• Round-3 blocker 1 (cross-variant guard also rejected VAE-only sourcing): fixed correctly. The split is right, every source-based encoder extraction goes through _validate_encoder_source, no VAE branch retains variant logic, and the linear UI can't wire a cross-variant source into an encoder slot in any main-format × standalone-VAE × standalone-encoder combination I enumerated.
• Round-3 blocker 2 (AutoProcessor-only local rungs): the offline gap itself is fixed — both local rungs now loop (AutoProcessor, AutoTokenizer). Unfortunately the widened AutoTokenizer rung introduces two new problems, one of which silently corrupts conditioning. Details below; both reproduced empirically against transformers 5.5.4 with real files, not hypothesized.
• The core_metadata mistral_encoder declaration + 2.2.0 bump looks good.

Housekeeping: I pushed two commits

• 7179712 — merge of current main. Conflicts: addRegions.ts (main refactored the posCond unions into a shared RegionalPositiveConditioning alias for Krea-2 regional prompting; I added flux2_dev_text_encoder to the alias and took main's shape), and schema.ts/openapi.json (regenerated from the merged backend rather than hand-merged).
• c2b7778 — the merge surfaced a semantic conflict git can't see: main's #9428 added a registry guard asserting every *_text_encoder node declares idle_gpu_offloadable=True, and flux2_dev_text_encoder (which neither parent knew about) failed it, so python-tests would have gone red. The flag alone would have been wrong — the marker's contract is that saved conditioning is CPU-backed, because the borrowed GPU's pool lock is released the moment the node returns. I mirrored the Klein encoder: detach().to("cpu") before save, marker added, version 1.0.1, plus the same output-device regression test Klein has. tests/app is fully green at head (2161 passed), as are the frontend suite (1736) and lints.

Blocker 1: AttributeError escapes the ladder and crashes the whole load

_TOKENIZER_LOAD_ERRORS = (OSError, EnvironmentError, ValueError, KeyError) (mistral_encoder.py:623) doesn't cover a failure mode of the very rung you added. In transformers 5.5.4, AutoTokenizer.from_pretrained on a directory whose tokenizer_config.json names a tokenizer_class unknown to 5.5.4 hits tokenizer_class_from_name(...) → None and raises AttributeError: 'NoneType' object has no attribute 'from_pretrained' (tokenization_auto.py:723 has no None-guard).

Triggering sequence (reproduced): encoder dir with config.json (model_type: "mistral3"), tekken.json, and a tokenizer_config.json containing "tokenizer_class": "MistralCommonTokenizer" — exactly what transformers 4.5x's MistralCommonTokenizer.save_pretrained writes. Rung 3's AutoProcessor raises OSError (caught), AutoTokenizer raises AttributeError → escapes _try_load_tokenizer_from_dir → _load_tokenizer_for_model crashes instead of falling through to the HF rung. Before this commit the same directory worked whenever rung 4 could succeed.

Given the ladder's whole contract is graceful fallthrough, I'd catch broadly in _try_load_tokenizer_from_dir (add AttributeError, or catch Exception and log) — an exotic exception from a probe rung should never be able to kill the load while a working fallback remains.

Blocker 2: root tekken.json + config.json now loads MistralCommonBackend, which tokenizes the template markers wrong — silently

The commit comment assumes tekken-carrying directories fail with KeyError: 'special_tokens'. That's only true when config.json is absent. With config.json (model_type: "mistral3") + a real tekken.json — the layout of an official mistralai Mistral-Small download, and the standalone-encoder layout these loaders explicitly accept — AutoTokenizer succeeds and returns MistralCommonBackend (mistral-common is a hard dependency, so this holds in every standard install).

That object BPE-encodes [SYSTEM_PROMPT]/[INST] as literal text instead of splicing them as single Tekken special ids. Verified against this module's own _TekkenRawTextAdapter on the same prompt: reference ends ..., 1, 17, 33153, 5117, 18, 3, 1097, 7990, 4 (markers = ids 17/18/3/4); MistralCommonBackend produces 1, 1091, 101289, 58343, ... (literal [, SY, STEM…). Every structural token is off-distribution, the encode "works", and generations silently degrade. Before this commit, that directory fell through to rung 4 and got the correct BFL processor whenever online/cached — so this regresses the online case too, and offline it trades an actionable RuntimeError for silent corruption.

Your ladder tests missed it because the fixture's tekken.json is a small fake that raises during parse; a real one parses fine and short-circuits.

Suggested fix: the ideal handling for a root tekken.json already exists in this file — wrap it with _TekkenRawTextAdapter (same as the embedded-tekken rung) before ever handing the directory to AutoTokenizer; and/or treat a mistral-common-backed result from AutoTokenizer as "not usable, keep falling". Either way, please add a fixture with a structurally valid tekken model file so the test exercises the success path, not just the raise path.

Non-blockers

1. The [dev]-side justification of the guard split overclaims, and the [dev] half is UI-dead. "buildFLUXGraph falls back to any FLUX.2 diffusers pipeline when only the VAE is needed" is true only of the Klein branch. The dev branch selects from dev-only pipelines (selectFlux2DevDiffusersModels) and readiness requires a dev diffusers source, so the config your test_dev_loader_extracts_vae_from_a_klein_pipeline proves the backend accepts (dev GGUF + standalone Mistral encoder + Klein diffusers as the only VAE source) still shows a disabled Invoke button. Safe direction, but either relax the dev builder/readiness VAE-sourcing to mirror Klein or trim the docstrings to the Klein case.
2. _validate_encoder_source in the Klein loader only rejects Dev, not Klein-family mismatches, while its comment claims the opaque-matmul class is closed. Klein 4B GGUF main + Klein 9B diffusers as qwen3_source_model (workflow editor) passes the guard and hits exactly that matmul error; the standalone-encoder path (_validate_qwen3_encoder_variant) and the frontend both enforce the family match — only the source path doesn't. Pre-existing, but worth closing while you're here: check the Klein→Qwen3 family map in the source path too. Related asymmetry: the dev guard rejects any non-dev variant, the Klein guard only dev — a future third FLUX.2 variant would be rejected on one side and silently accepted on the other.

Requesting changes for the two tokenizer-ladder blockers.

[lstein]

Round 5 review (head e74568d) — approved

Both round-4 blockers are fixed, and I verified each empirically by re-running my original repro directories through the real production code (not just the test fakes), plus a fresh-context adversarial pass over the commit.

• Blocker 1 (AttributeError escaping a probe rung): fixed. The probes now catch broadly with _TOKENIZER_LOAD_ERRORS demoted to log-level selection — the right shape: an unexpected failure is logged loudly instead of either crashing the load or being silently swallowed. The 4.5x-save_pretrained layout that used to kill the load now falls through cleanly (and in fact resolves correctly via the new tekken rung, since that layout carries a tekken.json).
• Blocker 2 (mistral-common tokenizer silently mis-encoding the template markers): fixed on both paths. The standalone-tekken.json rung ahead of the transformers probes handles the official mistralai layout, and _normalize_tokenizer catches anything that still arrives mistral-common-backed. I verified the re-wrap against the real transformers 5.5.4 MistralCommonBackend: the MRO module-name detection fires, .tokenizer is the underlying MistralTokenizer, and both paths produce token ids identical to the reference adapter — markers spliced as single Tekken ids. Re-wrapping instead of discarding was the right call for exactly the reason your commit message gives. I also confirmed the BFL FLUX.2-dev/tokenizer folder ships no tekken.json, so the new rung cannot shadow the canonical diffusers layout.
• Both non-blockers closed too. The Klein guard's allowlist on the shared _KLEIN_TO_QWEN3_VARIANT map covers all four Klein variants, fails closed for anything else (including a future variant), and shuts the 9B-source-into-4B-main workflow-editor gap; the hasattr→getattr fix in _validate_qwen3_encoder_variant and the corrected [dev] docstring are both right. Thanks also for the structurally valid tekken fixture — the ladder's success path is finally exercised, and I confirmed all five headline mutations (broad catch, tekken rung, re-wrap, family check, allowlist) each fail at least one new test.

CI is green at this head and the affected suites pass locally (944 tests).

Non-blocking residuals (fine as follow-ups or not at all)

1. _normalize_tokenizer only inspects the top-level object, so a processor whose inner .tokenizer is mistral-common-backed would pass through unnormalized. Unreachable today (any local dir with tekken.json is intercepted by the tekken rung first, and the BFL fallback resolves a fast tokenizer), and it behaved identically before this commit — noting it only so the assumption is written down.
2. _TekkenRawTextAdapter._resolve_special_ids returns {} when the wrapped tokenizer exposes neither accessor, and _encode then BPE-encodes the markers with no warning — the very failure mode this commit eliminates elsewhere. Pre-existing, but a one-line warning when _special_ids comes up empty would close the last silent path.

Approving. Thanks for the thorough turnaround across all five rounds — the tokenizer ladder is in much better shape than where this started.