| --- |
| license: mit |
| --- |
| |
| # 🍰 Tiny AutoEncoder for FLUX.2 |
|
|
| [TAEF2](https://github.com/madebyollin/taesd) is very tiny autoencoder which uses the same "latent API" as FLUX.2's VAE. |
| FLUX.2 is useful for real-time previewing of the FLUX.2 generation process, as well as general resource-constrained encoding/decoding. |
|
|
| This repo contains `.safetensors` versions of the TAEF2 weights. |
|
|
| ## Using in 🧨 diffusers |
|
|
| **NOTE**: Unlike TAEF1, TAEF2's architecture [isn't properly integrated](https://github.com/madebyollin/taesd/issues/35#issuecomment-3765620926) into Diffusers yet. |
| So for now you'll want some wrapper code: |
|
|
| ```sh |
| pip install git+https://www.github.com/huggingface/diffusers # needed for Klein support as of 2026-01-18 |
| wget -nc -nv https://raw.githubusercontent.com/madebyollin/taesd/refs/heads/main/taesd.py -O taesd.py |
| wget -nc -nv https://huggingface.co/madebyollin/taef2/resolve/main/taef2.safetensors -O taef2.safetensors |
| ``` |
|
|
| ```python |
| # 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 |
| ``` |
|
|
|
|
|  |
|
|
| ## Quality Comparisons |
|
|
| These compare TAEF2, the [full FLUX.2 VAE](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/blob/main/vae/config.json), and the alternate [FAL FLUX.2 Tiny](https://huggingface.co/fal/FLUX.2-Tiny-AutoEncoder) AE. |
|
|
|  |
|
|
|  |
|
|
|  |
|
|
|  |
