Feat: flux2 dev support

[Pfannkuchensack]

Summary

Adds end-to-end support for FLUX.2 [dev] alongside the existing FLUX.2 Klein implementation. Dev is a 32B guidance-distilled rectified flow transformer that uses the BFL "cow-mistral3-small" 30-layer Mistral distillation as its sole text encoder (Mistral Small 3 family, hidden 5120, joint_attention_dim 15360). It shares the 32-channel AutoencoderKLFlux2 VAE and the 4D-RoPE sampling backend with Klein, so most of the existing infrastructure is reused — only Mistral-specific loaders, configs, and graph wiring are new.

A key finding during this work: only the 30-layer cow distillation works. Upstream Mistral Small 3.1 / 3.2 (40 layers) produces off-distribution embeddings under FLUX.2's static (10, 20, 30) hidden-state extraction, because the joint attention was trained against the 30-layer model. The probing layer now rejects 40-layer Mistrals outright with a clear error.

Backend — taxonomy & probing

• Taxonomy: Flux2VariantType.Dev, ModelType.MistralEncoder, ModelFormat.MistralEncoder, single-variant MistralVariantType.Cow. Added to AnyVariant and variant_type_adapter. ModelRecordChanges.variant union extended.
• Probing: Dev is detected by context_in_dim = 15360 (main) / vec_in_dim = 5120 / hidden_size = 6144 (LoRA, all formats). Existing Klein and FLUX.1 probes preserved. Main_Diffusers_Flux2_Config also accepts Flux2Pipeline / Flux2Transformer2DModel.
• configs/mistral_encoder.py: Diffusers / single-file / GGUF configs under BaseModelType.Any / ModelType.MistralEncoder. All three formats reject non-cow Mistrals (hidden_size != 5120 or num_hidden_layers != 30) with a NotAMatchError that names the expected geometry. Diffusers folder probe excludes full pipelines (those match Main_Diffusers_Flux2_Config).

Backend — Mistral loaders
New load/model_loaders/mistral_encoder.py with three loaders (Diffusers via AutoModel, single-file via MistralModel, GGUF via MistralModel + llama.cpp key conversion). Includes:

• _convert_for_bare_mistral_model: strips the model. prefix and drops lm_head so a MistralForCausalLM state dict loads cleanly into bare MistralModel.
• _drop_quantization_metadata: dequantizes Comfy-Org's FP8/FP4 weights in place using the per-layer *.weight_scale (and optional *.input_scale) tensors and strips the _quantization_metadata / scaled_fp8 markers before load_state_dict.
• _materialize_remaining_meta_tensors: replaces any params/buffers still on the meta device after load_state_dict (norms → ones, others → zeros) with a warning listing what was missing, so the model cache → VRAM move can't fail on partial state dicts.
• llama.cpp converter handles attn_q_norm/attn_k_norm (Mistral 3.x QK-norm), attn_q/k/v/output, attn_norm, ffn_*, plus the root token_embd/output_norm/output keys.
• GGUF metadata reader pulls llama.rope.freq_base (1e9) and llama.context_length (131072) from the header so the rebuilt MistralConfig matches the trained-against rope.

Backend — embedded Tekken tokenizer
Both Comfy-Org's single-file safetensors and gguf-org's cow GGUFs ship the canonical Mistral Tekken tokenizer JSON embedded as a tensor named tekken_model — Comfy as a U8 blob, cow GGUFs as one fp16 value per byte. New _extract_tekken_bytes() reads either layout, _TekkenChatTemplateAdapter wraps mistral_common.MistralTokenizer to expose the HF apply_chat_template(tokenize=True, return_tensors='pt', padding='max_length', …) surface the invocation uses, and _load_tokenizer_for_model() walks: embedded Tekken → sibling diffusers tokenizer/ → black-forest-labs/FLUX.2-dev:tokenizer HF fallback. Adds mistral-common as a dependency.

This makes the Mistral encoder self-contained for the canonical Comfy / cow distributions — no HF tokenizer fetch needed.

Backend — invocations

• flux2_dev_model_loader, flux2_dev_text_encoder, flux2_dev_lora_loader (+ collection variant). MistralEncoderField added to model.py.
• Text encoder runs Mistral's chat template with a fixed system message (matches the diffusers _get_mistral_3_small_prompt_embeds reference) and stacks hidden states from layers (10, 20, 30) → (B, seq, 15360) matching the transformer's joint_attention_dim. For the 30-layer cow that's exactly (1/3, 2/3, last). Tries multimodal [{type, text}] content first (PixtralProcessor / Mistral3Processor / Tekken adapter) and falls back to plain-string content, then to manual [INST]…[/INST] formatting.
• flux2_denoise / flux2_vae_decode / flux2_vae_encode reused unchanged — already model-agnostic.

Qwen3 probe strictness (bugfix)

• _get_qwen3_variant_from_state_dict / _get_variant_from_config now return None / raise NotAMatchError for unknown hidden_size instead of silently defaulting to qwen3_4b. The old fallback meant any llama.cpp GGUF causal LM (Mistral, Llama, …) was misclassified as Qwen3 — caught when a Mistral 3.x GGUF was identified as qwen3_4b.

Frontend

• New MistralEncoderModelConfig type + isMistralEncoderModelConfig, isFlux2DevMainModelConfig, isFlux2DevDiffusersMainModelConfig guards. selectMistralEncoderModels, selectFlux2DevDiffusersModels, useMistralEncoderModels, useFlux2DevDiffusersModels hooks/selectors.
• paramsSlice: flux2DevVaeModel / flux2DevMistralEncoderModel / flux2DevSourceModel fields + reducers + selectIsFlux2Dev / selectIsFlux2Klein selectors.
• ParamFlux2DevModelSelect component (VAE + Mistral Encoder dropdowns), wired into AdvancedSettingsAccordion (Dev shows the Mistral selector, Klein keeps the Qwen3 selector).
• buildFLUXGraph: dev branch for flux2_dev_model_loader / flux2_dev_text_encoder / shared flux2_denoise with full txt2img / img2img / inpaint / outpaint support, plus multi-reference image editing via flux_kontext collect chain (Flux2RefImageExtension is model-agnostic). Dev model loader's vae is wired into both flux2_denoise and flux2_vae_decode.
• New addFlux2DevLoRAs helper, wires LoRAs through flux2_dev_lora_collection_loader.
• readiness.ts: variant-aware FLUX.2 readiness — dev requires flux2DevVaeModel + flux2DevMistralEncoderModel (or a Dev diffusers source), Klein keeps the Qwen3/VAE check. hasFlux2DevDiffusersSource threaded through both generate and canvas tabs.
• zModelType / zModelFormat / zFlux2VariantType extended for mistral_encoder / dev. zMistralVariantType enum carries only cow_mistral3_small. Display-name maps updated.
• New i18n keys: noFlux2DevVaeModelSelected, noFlux2DevMistralEncoderModelSelected, metadata.mistralEncoder.
• OpenAPI schema regenerated; TS types up to date.

Frontend — metadata recall

• New Flux2DevVAEModel + Flux2DevMistralEncoderModel handlers in parsing.tsx. Both gate on base === 'flux2' and use the presence of metadata.mistral_encoder (dev only) vs metadata.qwen3_encoder (Klein only) to disambiguate — dispatching into flux2DevVaeModelSelected / flux2DevMistralEncoderModelSelected instead of the Klein slices.
• KleinVAEModel updated to also reject when mistral_encoder is present in the metadata, so it no longer swallows Dev images into the Klein VAE slice.
• Wired both new handlers into the Recall Parameters panel in ImageMetadataActions.tsx (the panel iterates a hardcoded handler list; Klein VAE / Encoder were the only FLUX.2 entries before). Remix Image already iterates Object.values(ImageMetadataHandlers) automatically.

Starter models (cow-only Mistral encoders)

• black-forest-labs/FLUX.2-dev Diffusers (~80 GB, Non-Commercial)
• diffusers/FLUX.2-dev-bnb-4bit Diffusers (NF4)
• gguf-org/flux2-dev-gguf Q3_K_M / Q4_K_M / Q5_K_M / Q6_K / Q8_0 transformer-only entries depending on the FLUX.2 VAE + a cow Mistral encoder
• Mistral encoders (cow-only):
  • Comfy-Org single-file safetensors: bf16 (35.6 GB), fp8 (18 GB), fp4_mixed (12.3 GB) — all embed the Tekken tokenizer
  • gguf-org cow GGUFs: Q4_0, Q8_0, IQ4_XS — also embed Tekken
• Upstream Mistral Small 3.1 / 3.2 entries removed — they don't work for FLUX.2 (wrong layer count).

Related Issues / Discussions

None yet. The upstream feature/flux2-noncommercial-license work is unrelated but worth coordinating with — FLUX.2 [dev] inherits the BFL Non-Commercial License and is flagged in isNonCommercialMainModelConfig.

QA Instructions

Backend probing (no model load required, takes seconds):

uv run --extra cuda pytest tests/test_imports.py tests/model_identification/test_identification.py::test_default_settings_main tests/model_identification/test_identification.py::test_controlnet_t2i_default_settings

Manual probe against any local FLUX.2 [dev] artifacts (folder, transformer subfolder, text_encoder subfolder, GGUF, single-file VAE, cow Mistral). All should classify with the matching dev / mistral_encoder configs and existing FLUX.2 Klein fixtures should still classify identically. Probing a 40-layer Mistral Small 3.1 / 3.2 GGUF or safetensors should be rejected with a clear "expected hidden_size=5120, num_hidden_layers=30" error — this is intentional.

Frontend checks (from invokeai/frontend/web):

pnpm lint:tsc
pnpm lint:eslint
pnpm test:no-watch

End-to-end with a real model:

1. Install via the Model Manager:
   • Transformer: any of the gguf-org/flux2-dev-gguf quants (Q4_K_M is a good balance at ~20 GB) or the full black-forest-labs/FLUX.2-dev Diffusers folder
   • VAE: flux2-vae.safetensors (or use the one bundled in the Diffusers source)
   • Mistral encoder: either Comfy-Org/flux2-dev/split_files/text_encoders/mistral_3_small_flux2_fp8.safetensors (recommended — 18 GB, embedded Tekken) or any gguf-org/flux2-dev-gguf/cow-mistral3-small-q*.gguf variant.
2. Select the FLUX.2 [dev] transformer in the main-model picker. The Advanced Settings accordion should now show "FLUX.2 [dev] VAE" and "FLUX.2 [dev] Mistral Encoder" dropdowns (not the Klein Qwen3 selector).
3. For GGUF / single-file transformers: pick the standalone VAE + Mistral encoder. The pre-queue check should produce Dev-specific error messages if either is missing.
4. Defaults from MainModelDefaultSettings.from_base(Flux2, Dev): 28 steps, guidance 3.5, CFG 1.0, 1024×1024.
5. Generate. Watch the log on first run — for Comfy / cow files you should see Loaded embedded Tekken tokenizer from <filename> (no HF roundtrip). Subsequent runs go straight to denoise.
6. Try a multi-reference image: add a FLUX.2 ref image (the same UI as Klein), generate.
7. Apply a FLUX.2 LoRA and generate — variant mismatch (Klein LoRA on Dev) should log a warning but not crash.
8. Recall test: open any generated FLUX.2 [dev] image. Under "Recall Parameters" the panel should now include Mistral Encoder and VAE rows alongside Model / Steps / Scheduler. Click "Remix Image" — both Dev VAE and Mistral encoder should restore into the params slice. Repeat with a Klein image and confirm Klein VAE + Qwen3 encoder still recall correctly (no cross-contamination).

Embedded-Tekken caveat: if you install an older cow GGUF that doesn't embed tekken_model, or a diffusers folder without a sibling tokenizer/, the loader falls back to black-forest-labs/FLUX.2-dev:tokenizer on HF. With no internet and no HF cache, you'll see a clear error documenting three workarounds (install a Comfy-Org safetensors that bundles Tekken, populate the HF cache, or run huggingface-cli download black-forest-labs/FLUX.2-dev --include 'tokenizer/*').

Merge Plan

• No DB migration: the params slice gained new fields, but they default to null and the slice version did not change. Existing user states keep working — Dev-specific fields just stay null until the user picks a Dev model.
• Adds mistral-common to pyproject.toml dependencies (~6 MB + pycountry transitive ~8 MB). Required for the embedded Tekken tokenizer path; without it the loader still works via HF fallback.
• Tested locally end-to-end with flux2-dev-Q2_K.gguf transformer + Comfy fp8 Mistral encoder + standalone FLUX.2 VAE — full prompt adherence confirmed on the canonical "origami lighthouse on cliff" test prompt. Lower-fidelity setups (base Mistral 3.x at any quant level) produce visibly degraded compositions, which is why those configs are now rejected at probe time.
• Non-Commercial License: dev inherits the existing FLUX dev / Klein 9B non-commercial flagging in isNonCommercialMainModelConfig.

Checklist

• The PR has a short but descriptive title, suitable for a changelog — e.g. feat(flux2): add FLUX.2 [dev] support (cow Mistral, embedded Tekken, recall)
• Tests added / updated (if applicable) — model-identification + readiness fixtures cover new fields; metadata-recall test suite gained Flux2DevVAEModel + Flux2DevMistralEncoderModel cases including Klein/dev cross-rejection
• ❗Changes to a redux slice have a corresponding migration — N/A, new fields default to null, slice version unchanged
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)