fix(flux2-vae): support FLUX.2 small-decoder VAE variant

[Pfannkuchensack]

Summary

Since AutoencoderKLFlux2 only exposes a single block_out_channels for both encoder and decoder, the loader now:

1. Reads encoder.conv_in.weight / encoder.conv_norm_out.weight to derive the encoder channels.
2. Reads decoder.conv_in.weight / decoder.conv_norm_out.weight to derive the decoder channels.
3. Builds the model with the encoder's channels.
4. If the decoder differs, replaces model.decoder with a fresh Decoder instance using the smaller channel widths before load_state_dict is called.

The standard FLUX.2 VAE remains unaffected — when encoder and decoder match, no rebuild happens.

Not supported

The small_decoder.safetensors file from the same repository is not supported by this PR. It only contains the decoder (no encoder) and uses bare BFL keys without an encoder./decoder. prefix, so it neither passes the VAE probe heuristics nor can be loaded into a full AutoencoderKLFlux2 instance. Use full_encoder_small_decoder.safetensors instead, which bundles the original full encoder together with the small decoder.

Related Issues / Discussions

—

QA Instructions

1. Download full_encoder_small_decoder.safetensors from black-forest-labs/FLUX.2-small-decoder.
2. Install it as a FLUX.2 VAE via the model manager.
3. Run a FLUX.2 generation using this VAE and verify there is no size mismatch / expected shape error and the decoded image looks correct.
4. Regression check: load the standard FLUX.2 VAE and confirm it still works unchanged.

Merge Plan

Normal merge — change is isolated to Flux2VAELoader._load_model.

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]

Works as advertised.