tiny-flux-klien-3-vae

Tiny decoder for FLUX 3 Image Klein latents, fine-tuned from TAEF2 (MIT). FLUX 3 Klein ships the FLUX.2 VAE byte-for-byte (ae.safetensors sha256 868fe7b3...ce8f3 in both repos), so TAEF2's latent API matches; these weights are trained to decode it better.

• 2.68M params (1.34M encoder + 1.34M decoder) vs 84M for the full VAE; ~11x faster decode at 1024px.
• Works on the BN-normalized latent space (what the transformer outputs / what the diffusers pipeline denormalizes with vae.bn), exactly like TAEF2.
• Decoder was fine-tuned; the encoder is TAEF2's original encoder.
• Any size that is a multiple of 16 works (fully convolutional).

Files

file	what
tiny_klein3_v2.safetensors	recommended: long run (25.8k photos, L1 + LPIPS + light PatchGAN 0.01 + EMA 0.999, 512px then 1024px; best Kodak-LPIPS checkpoint)
tiny_klein3_ft.safetensors	first fine-tune (DIV2K, 3k steps, L1 + LPIPS)
checkpoints/	intermediate best/last checkpoints and the GAN discriminator (disc.pt) for resuming
code/	taesd.py (model def), training / eval / generation scripts
results/	eval grids, FLUX 3 generation comparisons, logs

Usage

import safetensors.torch as stt
from taesd import TAESD  # code/taesd.py

tiny = TAESD(encoder_path=None, decoder_path=None, latent_channels=32, arch_variant="flux_2")
tiny.load_state_dict(stt.load_file("tiny_klein3_v2.safetensors"))
# z_norm: BN-normalized, unpatchified latents (B, 32, H/8, W/8)
image_01 = tiny.decoder(z_norm).clamp(0, 1)

Drop-in for a diffusers pipeline: use the DiffusersTAEF2Wrapper from the TAEF2 model card and load this file instead of taef2.safetensors.

Results

Kodak (24 photos), reconstruction vs original:

reconstruction vs original image (mean over 24 images)
path                                PSNR dB   LPIPS
full VAE enc -> full VAE dec          34.38  0.0137
taef2.safetensors: full enc -> TINY dec    30.69  0.0379
tiny_klein3_ft.safetensors: full enc -> TINY dec    32.05  0.0358
tiny_klein3_v2.safetensors: full enc -> TINY dec    31.99  0.0306
taef2.safetensors: TINY enc -> full dec    31.75  0.0305
tiny_klein3_ft.safetensors: TINY enc -> full dec    31.75  0.0305
tiny_klein3_v2.safetensors: TINY enc -> full dec    31.75  0.0305
taef2.safetensors: TINY enc -> TINY dec    30.19  0.0405
tiny_klein3_ft.safetensors: TINY enc -> TINY dec    31.29  0.0400
tiny_klein3_v2.safetensors: TINY enc -> TINY dec    31.23  0.0356

decode latency (bf16, batch 1)
There are modules in AutoencoderKLFlux2 that should be kept in float32: []. Casting directly with `to()` can lead to inconsistent results; set `torch_dtype` in `from_pretrained()` instead to keep these modules in float32.
  1024px  full VAE    162.0 ms (2.7 GB peak) | taef2.safetensors   14.4 ms (0.9 GB) -> 11.2x faster | tiny_klein3_ft.safetensors   14.4 ms (0.9 GB) -> 11.2x faster | tiny_klein3_v2.safetensors   14.5 ms (0.9 GB) -> 11.2x faster
  2048px  full VAE    762.7 ms (9.8 GB peak) | taef2.safetensors   73.2 ms (2.8 GB) -> 10.4x faster | tiny_klein3_ft.safetensors   73.2 ms (2.8 GB) -> 10.4x faster | tiny_klein3_v2.safetensors   73.3 ms (2.8 GB) -> 10.4x faster

FLUX 3 Klein base generations (mixed sizes), vs full-VAE decode:

mean vs full-VAE decode of the same generated latents:
  taef2                PSNR 30.00 dB  LPIPS 0.0342
  tiny_klein3_ft       PSNR 31.47 dB  LPIPS 0.0265
  tiny_klein3_v2       PSNR 31.45 dB  LPIPS 0.0234

Write-up

https://tanmaypatil123.github.io/posts/2026/tiny-vae-flux3-klein

Notes

• Training data: DIV2K + Unsplash Lite (real photos only, no FLUX outputs). No images from either dataset are included in this repo; check the DIV2K and Unsplash dataset terms before reusing the data.
• Kodak was used both to pick the best checkpoint and to report numbers, so treat the Kodak numbers as slightly optimistic; the FLUX 3 generation numbers are independent.
• Not affiliated with or endorsed by Black Forest Labs.