feat(flux2): add FLUX.2 klein model support

[Pfannkuchensack]

Summary

Add initial support for FLUX 2 klein models architecture.

New Invocation Nodes

• flux2_denoise: Denoising invocation for FLUX 2 models
• flux2_klein_model_loader: Model loader for Klein architecture
• flux2_klein_text_encoder: Text encoder using Qwen3-based encoding
• flux2_vae_decode: VAE decoder for FLUX 2 output

Backend Support

• New flux2 module with denoise and sampling utilities
• Extended model manager configs for FLUX 2 model types
• Updated model loaders to support Klein architecture
• Added Qwen3 encoder configuration updates

Frontend Updates

• Extended graph builder for FLUX 2 workflow support
• Added FLUX 2 model types and configurations
• Updated readiness checks for FLUX 2 models
• UI components updated for new model type

Related Issues / Discussions

QA Instructions

1. Install a FLUX 2 Kontext model (Klein architecture)
2. Select the model in the UI
3. Test text-to-image generation with the new model
4. Verify step callbacks and progress display work correctly

Merge Plan

Standard merge after review.

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)

[lstein]

I tested with two models: the diffusers model at black-forest-labs/FLUX.2-klein-9B, and a quantized GGUF at https://huggingface.co/unsloth/FLUX.2-klein-4B-GGUF/resolve/main/flux-2-klein-4b-Q8_0.gguf

Initial observations:

1. When using the diffusers model and the built-in text encoder from the main model, I get the error:

ValueError: Cannot use chat template functions because tokenizer.chat_template is not set and no template argument was passed! For information about writing templates and setting the tokenizer.chat_temp
late attribute, please see the documentation at https://huggingface.co/docs/transformers/main/en/chat_templating                   

2. When using the diffusers model and clicking on the Qwen3 Encoder popup menu, I'm given the option to choose among a few previously-installed Z-Image Qwen3 Text Encoders (which may not be the desired behavior?) I choose Z-Image Qwen3 Text Encoder, but generation fails with

File "/home/lstein/Projects/InvokeAI/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py", line 82, in _autocast_forward                              
    return torch.nn.functional.linear(input, weight, bias)                                                                                                                                                
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                           
RuntimeError: mat1 and mat2 shapes cannot be multiplied (512x7680 and 12288x4096)

3. The GGUF file is recognized as a Flux1 model (not Flux.2).

4. In fact, there seems to be a regression with the Flux1 models as well. When using flux1 schnell or flux1 krea, I'm seeing this error on the front end: Failed to build graph: Node with id pos_cond_collect not found

[lstein]

Here's a trivial fix for the error Cannot use chat template functions because tokenizer.chat_template is not set...

diff --git a/invokeai/backend/model_manager/util/select_hf_files.py b/invokeai/backend/model_manager/util/select_hf_files.py
index 5bfb2f1f29..615eb394ef 100644
--- a/invokeai/backend/model_manager/util/select_hf_files.py
+++ b/invokeai/backend/model_manager/util/select_hf_files.py
@@ -70,6 +70,7 @@ def filter_files(
                 "weights.pb",
                 "onnx_data",
                 "spiece.model",  # Added for `black-forest-labs/FLUX.1-schnell`.
+                "chat_template.jinja", # added for `black-forest-labs/FLUX.2-klein-9B`.
             )
         ):
             paths.append(file)

[lstein]

The recent commit is causing the previously-installed ZImage Qwen3 encoders to fail validation with warnings like this one:

[2026-01-19 08:13:24,169]::[InvokeAI]::WARNING --> Skipping invalid model config in the database with name Z-Image Qwen3 Text Encoder. Ignoring this model. ({"key":"04198228-6d55-4c14-8759-069ae2eb4672","hash":"blake3:85a...)
[2026-01-19 08:13:24,169]::[InvokeAI]::WARNING --> Validation error: 1 validation error for tagged-union[Main_Diffusers_SD1_Config,Main_Diffusers_SD2_Config,Main_Diffusers_SDXL_Config,Main_Diffusers_SDXLRefiner_Config,Main_Diffusers_SD3_Config,Main_Diffusers_FLUX_Config,Main_Diffusers_Flux2_Config,Main_Diffusers_CogView4_Config,Main_Diffusers_ZImage_Config,Main_Checkpoint_SD1_Config,Main_Checkpoint_SD2_Config,Main_Checkpoint_SDXL_Config,Main_Checkpoint_SDXLRefiner_Config,Main_Checkpoint_FLUX_Config,Main_Checkpoint_Flux2_Config,Main_Checkpoint_ZImage_Config,Main_BnBNF4_FLUX_Config,Main_GGUF_FLUX_Config,Main_GGUF_ZImage_Config,VAE_Checkpoint_SD1_Config,VAE_Checkpoint_SD2_Config,VAE_Checkpoint_SDXL_Config,VAE_Checkpoint_FLUX_Config,VAE_Checkpoint_Flux2_Config,VAE_Diffusers_SD1_Config,VAE_Diffusers_SDXL_Config,ControlNet_Checkpoint_SD1_Config,ControlNet_Checkpoint_SD2_Config,ControlNet_Checkpoint_SDXL_Config,ControlNet_Checkpoint_FLUX_Config,ControlNet_Checkpoint_ZImage_Config,ControlNet_Diffusers_SD1_Config,ControlNet_Diffusers_SD2_Config,ControlNet_Diffusers_SDXL_Config,ControlNet_Diffusers_FLUX_Config,LoRA_LyCORIS_SD1_Config,LoRA_LyCORIS_SD2_Config,LoRA_LyCORIS_SDXL_Config,LoRA_LyCORIS_FLUX_Config,LoRA_LyCORIS_ZImage_Config,LoRA_OMI_SDXL_Config,LoRA_OMI_FLUX_Config,LoRA_Diffusers_SD1_Config,LoRA_Diffusers_SD2_Config,LoRA_Diffusers_SDXL_Config,LoRA_Diffusers_FLUX_Config,LoRA_Diffusers_ZImage_Config,ControlLoRA_LyCORIS_FLUX_Config,T5Encoder_T5Encoder_Config,T5Encoder_BnBLLMint8_Config,Qwen3Encoder_Qwen3Encoder_Config,Qwen3Encoder_Checkpoint_Config,Qwen3Encoder_GGUF_Config,TI_File_SD1_Config,TI_File_SD2_Config,TI_File_SDXL_Config,TI_Folder_SD1_Config,TI_Folder_SD2_Config,TI_Folder_SDXL_Config,IPAdapter_InvokeAI_SD1_Config,IPAdapter_InvokeAI_SD2_Config,IPAdapter_InvokeAI_SDXL_Config,IPAdapter_Checkpoint_SD1_Config,IPAdapter_Checkpoint_SD2_Config,IPAdapter_Checkpoint_SDXL_Config,IPAdapter_Checkpoint_FLUX_Config,T2IAdapter_Diffusers_SD1_Config,T2IAdapter_Diffusers_SDXL_Config,Spandrel_Checkpoint_Config,CLIPEmbed_Diffusers_G_Config,CLIPEmbed_Diffusers_L_Config,CLIPVision_Diffusers_Config,SigLIP_Diffusers_Config,FLUXRedux_Checkpoint_Config,LlavaOnevision_Diffusers_Config,Unknown_Config]
`qwen3_encoder.qwen3_encoder.any`.variant
  Field required [type=missing, input_value={'key': '04198228-6d55-4c...ormat': 'qwen3_encoder'}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.12/v/missing

[lstein]

@Pfannkuchensack Fantastic work! I am doing testing now. Hope to be able to approve and merge Friday.

[lstein]

@Pfannkuchensack With this PR I'm getting config validation errors on previously-installed Z-Image Turbo models. I can reinstall the models and they work. Is this expected?

[Pfannkuchensack]

@Pfannkuchensack With this PR I'm getting config validation errors on previously-installed Z-Image Turbo models. I can reinstall the models and they work. Is this expected?

I dont have a z-image installed in my test instance, so i cant check that. Do you have message from the console ?

[lstein]

@Pfannkuchensack With this PR I'm getting config validation errors on previously-installed Z-Image Turbo models. I can reinstall the models and they work. Is this expected?

I dont have a z-image installed in my test instance, so i cant check that. Do you have message from the console ?
False alarm. The validation error has nothing to do with FLUX.2 klein support. It is due to previous testing with my own run-on-CPU feature, which added a new field to the stored config for several models I was testing with.

So far the only issue I've encountered with your PR is that the VAE and text encoders (in the "Advanced" section) are not being recalled when you recall generation parameters.

Inpainting is working, but I find that I can't significantly modify the raster layer unless I set the denoising level to a high level of 0.95 or greater. Is this expected for this model?

I'm continuing to test various features.

[Pfannkuchensack]

The inpainting should be better now.

[lstein]

@Pfannkuchensack I think this is nearly ready to go. The only things missing for release are:

• Flux.2 model starter pack.
• Write a What's New item for the Flux.2 models (in en.json)
• The hardware requirements in the quick_start.md file should either split out FLUX.1 and FLUX.2 requirements, or note that they both have the same requirements (if that is the case).

A couple of questions for the release notes:

1. Are there any gotchas that users should be aware of?
2. Is there an ETA for controlnet and regional guidance?
3. Is there an ETA for the image edit functionality?

[Pfannkuchensack]

Added the image edit functionality. I did not test it yet but i saw that you can use openpose images as ref images and changes poses with klein

Prompt format may be important

According to diffusers documentation, you should reference images in your prompt:

• "the red car from image 1 changed to yellow"
• "change the color of the car to yellow" (simple style)
• For multiple images: "image 1", "image 2" etc.

FLUX.2 Klein only

• Only works with FLUX.2 Klein models (4B/9B)
• FLUX.1 still requires a separate Kontext model

No model selector

• Unlike FLUX.1, there is no dropdown for reference image models
• This is intentional - FLUX.2 Klein has this capability built-in

VRAM usage

• Each reference image increases the sequence length during denoising
• More reference images = more VRAM

[lstein]

Great! I'll do final review tonight and hopefully can merge after that. I'll take a look at VRAM consumption as edit images are added and add some guidance to the release notes.

[lstein]

I have tested all features, reviewed documentation, and everything checks out. Amazing work!

[foxhoundunit]

For some reason I can't use it. But it works fine in Comfy on my 4070 12Gb.
Also I can't change T5/Clip in Advanced tab.

[lstein]

For some reason I can't use it. But it works fine in Comfy on my 4070 12Gb. Also I can't change T5/Clip in Advanced tab.

Can you check the model manager to see what the problematic model has been autodetected as? It looks like it thinks that this is a flux.1 model, not flux.2.

[Pfannkuchensack]

For some reason I can't use it. But it works fine in Comfy on my 4070 12Gb.
Also I can't change T5/Clip in Advanced tab.

Did you build the frontend from the state of the PR?

[foxhoundunit]

For some reason I can't use it. But it works fine in Comfy on my 4070 12Gb. Also I can't change T5/Clip in Advanced tab.

Can you check the model manager to see what the problematic model has been autodetected as? It looks like it thinks that this is a flux.1 model, not flux.2.

For some reason I can't use it. But it works fine in Comfy on my 4070 12Gb.
Also I can't change T5/Clip in Advanced tab.

Did you build the frontend from the state of the PR?

I don't think so. I just installed a safetensors model file via Scan Folder tab with default settings

--

Ok I just updated to 6.11rc and downloaded klein 9bfp8 model from Starter pack, now it initializes correct.

But it still doesn't work. My qwen_3_8b_fp8mixed (https://huggingface.co/Comfy-Org/vae-text-encorder-for-flux-klein-9b/tree/main/split_files/text_encoders) defines as checkpoint and an error occurs during generation. As you can see Qwen3 4B text encoder from starter pack doesn't even shown in Advanced tab

This error occurs if I change my qwen_3_8b_fp8mixed model format to Qwen3 Encoder

[foxhoundunit]

By default a heavy 8B 16GB encoder is starts downloading when I select klein 9bfp8 from Starter pack. And I can even see it in an Advanced tab, but it's impossible for my 4070 12Gb / 64Gb

[Makhaon]

I've try to use new flux2 model on my laptop with NVIDIA GeForce RTX 5060 Laptop GPU, 8 Gb VRAM, but have no success. OOM occured. I've downloaded FLUX.2 Klein 4B (FP8) model, Invoke AI have installed FLUX.2 Klein Qwen3 4B Encoder
with this model. But it could not been load with 8 Gb VRAM:

Loading checkpoint shards: 100%|█████████████████████████████████████████| 2/2 [00:00<00:00, 66.62it/s]
[2026-01-29 13:25:51,876]::[ModelManagerService]::INFO --> [MODEL CACHE] Locking model cache entry db352
012-d90c-4a7c-9f8b-4c00996be98d:text_encoder (Type: Qwen3ForCausalLM), but it has already been dropped f
rom the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code (Se
e #7513).
[2026-01-29 13:25:54,978]::[ModelManagerService]::WARNING --> [MODEL CACHE] Insufficient GPU memory to l
oad model. Aborting
[2026-01-29 13:25:54,978]::[ModelManagerService]::WARNING --> [MODEL CACHE] Insufficient GPU memory to l
oad model. Aborting

I'm using v6.11.0.rc1 version of the programm. Does it possible to use this model in my situation?

As stated in the release documentation, this config should work:

FLUX.2 Klein 4B - 1024×1024
GPU: Nvidia 30xx series or later, 12GB+ VRAM (e.g. RTX 3090, RTX 4070). FP8 version works with 8GB+ VRAM.*

[Makhaon]

I just tryied 'FLUX.2 Klein 4B GGUF Q4_K_M quantized', that labeled 'runs on 6-8GB VRAM' with same result

[Pfannkuchensack]

Do you have enabled low vram mode https://invoke-ai.github.io/InvokeAI/features/low-vram/ ?

[Makhaon]

Hmm, I actually changed these parameters, but apparently a program update reset this file. I re-added the parameters and it helped, thank you very much!