[Feature] Flux.2[Dev] Support

[stduhpf]

Feature Summary

Support for the lastest Flux Image generation and editing model

Detailed Description

FLUX.2 [dev] is a 32 billion parameter rectified flow transformer capable of generating, editing and combining images based on text instructions. It uses a Mistral VLM (maybe pixtral? it's Mistral Small 3.2) as text encoder.

At first glance the architecture seems similar to Flux.1, with Double Blocks followed by Single Blocks
There are significant architecture change to the blocks themselves. I feel like it's going to need a separate implementation to Flux.1.

It uses a 32 Channel VAE (maybe 128 channels according to ComfyUI? I'm confused here).

Additional context

Announcement: https://bfl.ai/blog/flux-2

Model weights: https://huggingface.co/black-forest-labs/FLUX.2-dev

Inference code (Reference): https://github.com/black-forest-labs/flux2
Inference code (ComfyUI): comfyanonymous/ComfyUI#10879

Seems like a (smaller?) distilled version called Flux.2[klein] is also soon-to-be released.

[Green-Sky]

I have high hopes for the klein distill.

Looks like the dev version can only use about 6 reference images (as per their prompt guide)

Multi-Reference | Composite images | [pro]: 8, [flex]: 10, [dev]: ~6

[stduhpf]

Looks like the dev version can only use about 6 reference images (as per their prompt guide)

It probably means it only got trained with up to 6 references, but hopefully it's somewhat able to generalize to more references?

[stduhpf]

The VAE seems to be definitely 32 Channels, but there must be more to it, because If I just try to encode and then immediately decode an image I get broken results:

Edit: must be related to this:

[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.bn.running_mean | f32 | 1 [128, 1, 1, 1, 1]' in model file
[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.bn.running_var | f32 | 1 [128, 1, 1, 1, 1]' in model file
[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.post_quant_conv.bias | f32 | 1 [32, 1, 1, 1, 1]' in model file
[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.post_quant_conv.weight | f32 | 4 [1, 1, 32, 32, 1]' in model file
[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.quant_conv.bias | f32 | 1 [64, 1, 1, 1, 1]' in model file
[INFO ] model.cpp:1598 - unknown tensor 'first_stage_model.quant_conv.weight | f32 | 4 [1, 1, 64, 64, 1]' in model file

Yes it is:

            dd_config.z_channels = 32;
            use_quant = true;
            embed_dim = 32;

With that I could already start training the proj matrix for previews.

Still missing the bn stuff.

[leejet]

WIP!

[stduhpf]

I tried implementing the batch_norm thing, I'm not sure if it was worth the effort

Here's the difference when doing a round trip encode/decode using Flux.2 VAE with our without batch_norm:

original	No BatchNorm	With BatchNorm
difference to original	max: 76, min: -66, mean: 0.1725, std: 4.1707	max: 98, min: -124, mean: 0.4240, std: 5.8606

@leejet Let me know if you want me to provide the patch for implementing it.

Edit: Maybe my implementation is bad, but it looks like the results are actually worse with batchnorm? I consistently get more deviation from the original with batchnorm enabled... It seems very much not worth it.

[leejet]

Flux2 support: #1016.