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]

4B is working for me, although I find the results pretty underwhelming relative to the hype.

However, 9B is producing bad images. I did fresh installs of the official diffusers versions from HuggingFace using the MM interface. Prompt was "photograph of ginger cat, sitting on the sofa in a modern sunlit apartment. A tag around the cat's neck reads "'Klein'". I used CFG 1, the Euler scheduler, and 10 steps.

4B with 10 steps:

9B with 10 steps:

4B with 8 steps

9B with 8 steps

4B with 4 steps

9B with 4 steps

And just for fun, 9B with 8 steps and CFG=4

[lstein]

Another issue I just noticed. When I try to generate with Flux.1 (schnell or krea), I am getting "Node with id pos_cond_collect is not found."

[Edit]: I think this was a one-time mismatch between the front and back end; I'll report if it happens again.

[lstein]

@Pfannkuchensack Generations are looking good now. If you are satisfied with the image quality, then I think we can make a point release as soon as the following remaining issues are addressed:

• The standalone model (e.g. flux-2-klein-4b.safetensors) cannot automatically find its Qwen3 encoder in the main model
• Define FLUX.2 Klein starter models (including a compatible standalone encoder)
• Define default generation parameters for the klein models (steps, CFG, etc)
• Fix pydantic validation errors for the ZiT qwen3 encoder models
• Add FLUX.2 model requirements (minimum hardware needed) to the docs

Please also give me an idea of when these missing features may be added, so that I can address user expectations in the release notes:

• img2img
• inpainting
• outpainting
• support for quants / GGUF

Is the image editing feature accessed through img2img from the raster layer, or is there an additional head for the reference image?

[Pfannkuchensack]

Is the image editing feature accessed through img2img from the raster layer, or is there an additional head for the reference image?

I'll take a look at it later.

Add FLUX.2 model requirements (minimum hardware needed) to the docs

I'll do it later, or someone else can take care of it.

I did some testing with inpaiting, outpaiting and both together. i2i works fine.

[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 ?