Feat: flux2 dev support#9234
Open
Pfannkuchensack wants to merge 21 commits into
Open
Conversation
Adds end-to-end support for FLUX.2 [dev] alongside the existing Klein implementation. Dev uses Mistral Small 3.1 (24B) as its sole text encoder instead of Klein's Qwen3, with joint_attention_dim=15360 and the guidance-distilled 32B transformer. Backend - taxonomy: Flux2VariantType.Dev, ModelType.MistralEncoder, ModelFormat.MistralEncoder, MistralVariantType - configs: probe dev via context_in_dim=15360 (main + LoRA); new mistral_encoder.py with Diffusers / Checkpoint / GGUF configs; Main_Diffusers_Flux2_Config accepts Flux2Pipeline class name - loaders: new mistral_encoder.py (AutoModel for Diffusers folder, MistralModel for single-file + GGUF with llama.cpp key conversion). Existing Klein transformer loaders are generic enough for dev - ModelRecordChanges.variant union extended with MistralVariantType Invocations - flux2_dev_model_loader, flux2_dev_text_encoder (Mistral chat-template with FLUX2_DEV_SYSTEM_MESSAGE and layer-stacking 10/20/30), flux2_dev_lora_loader (+ collection variant) - MistralEncoderField on model.py; flux2_denoise / flux2_vae_decode / flux2_vae_encode reused unchanged (already model-agnostic) Frontend - types/hooks/selectors for MistralEncoder, isFlux2DevMainModelConfig, selectFlux2DevDiffusersModels, useMistralEncoderModels - params slice fields flux2DevVaeModel / flux2DevMistralEncoderModel / flux2DevSourceModel + reducers, selectIsFlux2Dev / selectIsFlux2Klein - ParamFlux2DevModelSelect component, wired into AdvancedSettingsAccordion - buildFLUXGraph dev branch with full txt2img / img2img / inpaint / outpaint + multi-reference image editing (same flux_kontext + collect chain as Klein, since Flux2RefImageExtension is model-agnostic) - addFlux2DevLoRAs helper for dev LoRA wiring - zModelType / zModelFormat / zFlux2VariantType extended for mistral_encoder / mistral_small_3_1 / dev - OpenAPI schema regenerated, TS types updated Starter models - FLUX.2 [dev] Diffusers (bf16 + NF4), three GGUFs (Q4/Q6/Q8), Mistral encoder (bf16 + NF4)
Follow-up fixes after first end-to-end run with FLUX.2 [dev] GGUF +
Mistral 3.x GGUF + standalone FLUX.2 VAE.
Frontend
- buildFLUXGraph: wire dev model loader's vae into both flux2_denoise
(required for BN statistics / inpaint) and flux2_vae_decode; missing
edge was raising RequiredConnectionException at runtime
- readiness.ts: variant-aware FLUX.2 readiness check — dev requires
flux2DevVaeModel + flux2DevMistralEncoderModel (or a Dev diffusers
source); Klein keeps Qwen3/VAE check. Threads
hasFlux2DevDiffusersSource through generate + canvas tabs and updates
buildGenerateTabArg / buildCanvasTabArg test helpers
- en.json: noFlux2DevVaeModelSelected, noFlux2DevMistralEncoderModelSelected
Mistral encoder loader (GGUF / single-file)
- Fix "Cannot copy out of meta tensor": llama.cpp conversion produced
`model.*` keys but loader instantiated bare MistralModel (no `model.`
prefix). Add _convert_for_bare_mistral_model to strip the prefix and
drop lm_head before load_state_dict
- _materialize_remaining_meta_tensors: after load_state_dict, replace any
still-meta parameters (norms→ones, others→zeros) and buffers so the
cache→VRAM move can't fail on partial state dicts, with a warning
listing what was missing
- llama.cpp converter: map attn_q_norm/attn_k_norm (Mistral 3.x qk-norm
variants), with ordering before attn_q/attn_k to avoid bad rewrites
Tokenizer / processor fallback
- _load_processor_with_offline_fallback walks a list of sources
(black-forest-labs/FLUX.2-dev tokenizer subfolder, then
mistralai/Mistral-Small-3.1-… and 3.2-…), trying AutoProcessor then
AutoTokenizer for each, cache-first then online. Final error spells
out the three workarounds (install Diffusers folder, set HF_ENDPOINT,
pre-cache the tokenizer)
- flux2_dev_text_encoder: try multimodal `[{type, text}]` chat template
first (PixtralProcessor / Mistral3Processor), fall back to plain
string content (AutoTokenizer), then to manual [INST]…[/INST]
Qwen3 encoder probe strictness
- _get_qwen3_variant_from_state_dict and _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 wrongly classified as Qwen3 —
visible when the Mistral 3.x GGUF was identified as a Qwen3-4B encoder
- Checkpoint / GGUF / Diffusers loaders propagate the strictness
Pfannkuchensack
requested review from
JPPhoto,
blessedcoolant,
dunkeroni and
lstein
as code owners
May 25, 2026 20:15
github-actions
Bot
added
python
…andlers Upstream Mistral Small 3.1/3.2 (40 layers) produces off-distribution embeddings under FLUX.2's static (10, 20, 30) hidden-state extraction. The joint attention was actually trained against BFL's 30-layer cow-mistral3-small distillation — both Comfy-Org's safetensors and gguf-org's cow GGUFs ship the same 30-layer weights, just packaged differently. - Probing (configs/mistral_encoder.py) now rejects non-cow Mistrals across all three formats (Diffusers / Checkpoint / GGUF) with a clear error. - Loader (load/model_loaders/mistral_encoder.py) extracts the embedded Tekken tokenizer from the `tekken_model` U8 (safetensors) / fp16-per-byte (cow GGUF) tensor via mistral_common, falling back to the BFL HF tokenizer. Removes the INVOKEAI_MISTRAL_TOKENIZER_SOURCE env var. - Starter models: drop upstream Mistral 3.x entries, add Comfy-Org bf16/fp8/fp4 variants alongside the cow GGUFs. - MistralVariantType: drop Small3_1, keep only Cow. - pyproject.toml: add mistral-common dependency. Frontend recall: - Add Flux2DevVAEModel + Flux2DevMistralEncoderModel handlers, disambiguating Klein vs dev via presence of `mistral_encoder` / `qwen3_encoder` metadata fields (both bases are `flux2`). - Wire both into the Recall Parameters panel (hardcoded list was missing them). - Add `metadata.mistralEncoder` i18n key + colocated tests.
…encoders
After studying ComfyUI's `Flux2Tokenizer` / `Mistral3_24BModel` reference
implementation, align the FLUX.2 [dev] text-encoder path with their setup:
- Probing now accepts both 30-layer (cow distillation) and 40-layer (Mistral
Small 3, BFL canonical / upstream) Mistrals. Re-adds `MistralVariantType.Mistral24B`
alongside `Cow`. All three configs (Diffusers / Checkpoint / GGUF) updated.
- Loaders strip `model.norm` (replace with Identity) when the loaded weights
are the 30-layer cow distillation. Matches Comfy's `final_norm=False` for
the pruned variant; for transformers' `MistralModel` the final RMSNorm is
always built but the cow was trained against the raw post-layer-29 state.
- 40-layer loads now log a clear warning that upstream Mistral 3.1 / 3.2 is
NOT what FLUX.2's joint attention was trained against and recommends the
Comfy-Org bf16/fp8/fp4 or gguf-org cow GGUF variants. BFL's canonical
bundled text_encoder is also 40-layer so we don't hard-reject; the warning
is opt-in self-discipline.
- Text encoder invocation switches from `apply_chat_template(messages, ...)`
to a raw text template `[SYSTEM_PROMPT]{sys}[/SYSTEM_PROMPT][INST]{prompt}[/INST]`
fed straight to the tokenizer — byte-for-byte matches Comfy's
`Flux2Tokenizer.llama_template.format(text)`. System prompt now includes
the literal `\n` between "object" and "attribution" Comfy ships.
- `_TekkenChatTemplateAdapter` renamed to `_TekkenRawTextAdapter` and exposes
a `__call__(text, padding_side='left', ...)` interface that Tekken-encodes
the raw string (BOS=1, no EOS) and left-pads with token id 11. Matches
Comfy's `pad_left=True` / `pad_token=11` settings.
Frontend types extended for the new `mistral3_24b` variant
(zMistralVariantType, MODEL_VARIANT_TO_LONG_NAME, schema.ts).
lstein
requested changes
Jul 10, 2026
Knip reported 6 unused exports. Each was dead code rather than incomplete wiring, verified against the actual consumers: - Drop the vestigial `flux2DevSourceModel` param end-to-end (state field, default, migration, reducer, action, selector, test). The FLUX graph builder auto-picks the diffusers source itself and never read this param; no UI set it. Mirrors how the Klein path already works. - Delete `selectIsFlux2Klein`; the graph builder computes this locally and only `selectIsFlux2Dev` is consumed. - Un-export `zMistralVariantType`; used only in the local `zAnyModelVariant` union, like `zQwenImageVariantType`. - Delete `selectMistralEncoderModels`; components use the `useMistralEncoderModels` hook instead. - Un-export `isFlux2DevMainModelConfig`; used only within types.ts, like its `isFluxDevMainModelConfig` / `isFlux2Klein9BMainModelConfig` siblings.
The dev-vs-Klein VAE recall keyed off the presence of a mistral_encoder metadata field, but that field is only written when a standalone Mistral encoder is selected. A FLUX.2 [dev] image whose encoder came from a Diffusers source has a vae field but no mistral_encoder, so its VAE was silently recalled into the Klein slice. Resolve the image's own main model and check variant === 'dev' instead — the same signal the graph builder uses. Add regression coverage for the mistral_encoder-absent dev case, and add the missing modelManager.flux2Dev* i18n keys so the [dev] VAE/encoder labels are translatable.
The diffusers FLUX.2-dev text encoder loads a PixtralProcessor, whose first positional __call__ parameter is `images`, not `text`. Passing the prompt positionally routed it into `images`, breaking the diffusers encoder path (only single-file/GGUF encoders, which use a text-first adapter, had been exercised). Pass text= explicitly.
The diffusers FLUX.2-dev text encoder loads a PixtralProcessor, whose first positional __call__ parameter is `images`, not `text`. Passing the prompt positionally raised "Incorrect image source", breaking the diffusers encoder path entirely. Only single-file/GGUF encoders (text-first adapter) had been exercised. Verified against transformers 5.5.4. fix(flux2): emit Tekken special tokens in the embedded-tokenizer adapter _TekkenRawTextAdapter used mistral_common's raw Tekkenizer.encode, which runs with SpecialTokenPolicy.IGNORE and BPE-encodes the FLUX.2 markers ([SYSTEM_PROMPT], [/SYSTEM_PROMPT], [INST], [/INST]) as literal text — 54 tokens instead of 36, corrupting the prompt structure fed to FLUX.2 on the single-file and GGUF paths. Resolve the marker ids from the tokenizer's special vocab and splice them in; output is now byte-identical to the reference PixtralProcessor.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_dim15360, hidden states sampled at layers 10/20/30) instead of Klein's Qwen3. It shares the 32-channelAutoencoderKLFlux2VAE 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-devdiffusers repo, and the 30-layercow-mistral3-smalldistillation (Comfy-Org safetensors, gguf-org GGUFs). Single-file and GGUF encoders embed the Tekken tokenizer, so no separate tokenizer download is needed.What works
Limitations / notes
tokenizer/) fall back to fetchingblack-forest-labs/FLUX.2-dev:tokenizerfrom HF; offline with no HF cache this errors with documented workarounds.isNonCommercialMainModelConfig).How to test
Automated:
# backend (fast, no model load) uv run --extra cuda pytest tests/test_imports.py tests/model_identification/test_identification.pypyproject.tomlgainedmistral-common;uv.lockis regenerated (uv lock --lockedpasses).End-to-end:
gguf-org/flux2-dev-ggufquant — ~20 GB for Q4_K_M — or the fullblack-forest-labs/FLUX.2-devdiffusers folder), the FLUX.2 VAE, and a Mistral encoder (Comfy-Org fp8 recommended, or a Q6/Q8 cow GGUF).Loaded embedded Tekken tokenizer from <file>(no HF fetch).Checklist
What's Newcopy (if doing a release after this PR)Closes #8668