TAEF2

TAEF2

Previews FLUX.2 image generation in real time with a tiny autoencoder that shares the FLUX.2 VAE's latent format. For image app developers and people with limited memory.

Made by madebyollin on Hugging Face. These notes are Mindket's, written from its model page and model card; the model's own files and license are on Hugging Face.

About

What it is: TAEF2 is a tiny autoencoder for FLUX.2, made by madebyollin. It speaks the same latent format as the FLUX.2 VAE, so it can stand in for it. The card names two uses: real-time previews while FLUX.2 generates, and encoding or decoding when resources are tight. This repository holds the weights as safetensors files.

Who it suits: It suits developers who build image generation tools on FLUX.2 and people who run it on limited hardware.

Industry | Who uses it | How they use it | Example
Software | Image app developers | Show a live preview while FLUX.2 generates | Preview a render before it finishes
Hobby | People with small GPUs | Encode and decode latents with less memory | Run FLUX.2 Klein with CPU offloading
Research | Model researchers | Compare tiny and full autoencoders | Check TAEF2 against the full FLUX.2 VAE

People who only want to make images with a ready app can skip it, since this is a component for pipelines.

Where to use it:

• Python with the diffusers library, through the wrapper code the card gives
• A FLUX.2 Klein pipeline, where it replaces the VAE
• Windows, macOS or Linux with Python, for downloading the files

Use case | Who | What they get back
Live previews | App developers | A quick decoded image while generation runs
Low-memory decoding | Hobbyists | Images decoded with a much smaller model
Comparison tests | Researchers | Side-by-side results against two other autoencoders

How to set it up: You need a few things first.

• Python with torch and safetensors
• A diffusers build that supports FLUX.2 Klein, installed from source
• The taesd.py file from the taesd project, which the wrapper imports
• The taef2.safetensors weights from this repository

1. Install a diffusers build with Klein support.
2. Place taesd.py next to your script.
3. Download the taef2.safetensors weights.
4. Paste the wrapper code from the card into your script.

pip install -U huggingface_hub
hf download madebyollin/taef2

How to use it: TAEF2 is not properly integrated into diffusers yet, so you wrap it in a small class that looks like a diffusers VAE. Then you assign it to the pipeline's vae and generate as usual.

# Construction
from taesd import TAESD
import torch
import safetensors.torch as stt
from diffusers.utils.accelerate_utils import apply_forward_hook

def convert_diffusers_sd_to_taesd(sd):
    out = {}
    for k, v in sd.items():
        encdec, _layers, index, *suffix = k.split(".")
        offset = 0
        if encdec == "decoder":
            offset = +1
        out[".".join([encdec, str(int(index)+offset), *suffix])] = v
    return out

class DotDict(dict):
    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__

class DiffusersTAEF2Wrapper(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dtype = torch.bfloat16
        self.taesd = TAESD(encoder_path=None, decoder_path=None, latent_channels=32, arch_variant="flux_2").to(self.dtype)
        self.taesd.load_state_dict(convert_diffusers_sd_to_taesd(stt.load_file("taef2.safetensors")))
        self.bn = torch.nn.BatchNorm2d(128, affine=False, eps=0.0) # default bn
        self.config = DotDict(batch_norm_eps=self.bn.eps)

    @apply_forward_hook
    def encode(self, x):
        return DotDict(latent_dist=DotDict(sample=lambda : self.taesd.encoder(x.to(self.dtype).mul(0.5).add_(0.5)).to(x.dtype)))

    @apply_forward_hook
    def decode(self, x, return_dict=True):
        x = self.taesd.decoder(x.to(self.dtype)).mul(2).sub_(1).clamp_(-1, 1).to(x.dtype)
        return dict(sample=x) if return_dict else x,

taef2_diffusers = DiffusersTAEF2Wrapper().eval().requires_grad_(False)

# Usage
from diffusers import Flux2KleinPipeline

device = "cuda"
dtype = torch.bfloat16

pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=dtype)
pipe.vae = taef2_diffusers
pipe.enable_sequential_cpu_offload() # pipe.enable_model_cpu_offload() # pipe = pipe.to(device)

prompt = "A slice of delicious New York-style berry cheesecake"
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=1.0,
    num_inference_steps=4,
    generator=torch.Generator(device="cpu").manual_seed(0)
).images[0]
image.save("flux-klein.png")
image

Tool | What it does
taesd.py | Provides the TAESD class that the wrapper builds on
taef2.safetensors | Holds the TAEF2 weights
DiffusersTAEF2Wrapper | Makes TAEF2 look like a diffusers VAE
Flux2KleinPipeline | The FLUX.2 Klein pipeline whose vae you replace
enable_sequential_cpu_offload | Moves parts of the model to the CPU to save GPU memory

Example: You ask the pipeline for a slice of New York-style berry cheesecake at 1024 by 1024 with 4 steps. The decoder you swapped in produces the picture, and the script saves it as flux-klein.png.

For newcomers and experts: An autoencoder converts between pictures and the compact latent form that image models work in. If you are new, run the card's example as written and look at flux-klein.png.

Because TAEF2 keeps the same latent interface as the FLUX.2 VAE, you can swap it into a pipeline without changing anything else. The wrapper shifts decoder layer indices when it converts the diffusers weights, uses 32 latent channels with the flux_2 architecture variant, and runs in BF16.

FLUX.2 pipeline produces latents
Wrapper passes them to the TAEF2 decoder
Decoder returns pixels scaled to minus 1 to 1
Image is shown or saved as a file

Good to know:

• License: MIT.
• Cost to run: no API key or paid service named. The example downloads FLUX.2 Klein 4B from Hugging Face and uses a CUDA device.
• Data it reaches: the latents and images your pipeline produces. The card names no account or web access.
• Not for: use in diffusers without the wrapper, since the card says it is not properly integrated there yet.

On Hugging Face, 13 November 2026

• Downloads in the last 30 days: 0
• Likes: 32
• License: mit
• Quantized versions: 1

Links

• Model: https://huggingface.co/madebyollin/taef2
• Files: https://huggingface.co/madebyollin/taef2/tree/main

Get the model

Download its files with Hugging Face's command line:

pip install -U huggingface_hub
hf download madebyollin/taef2

Then follow the usage in its model card. Nothing in this kit runs by itself.