File size: 3,874 Bytes
c6edc9d 21cb502 c6edc9d 21cb502 c6edc9d 901402f c6edc9d 21cb502 c6edc9d 901402f c6edc9d 7a18a47 c6edc9d 21cb502 d01e1b3 bd244eb d01e1b3 bd244eb d01e1b3 bd244eb d01e1b3 bd244eb d01e1b3 bd244eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | ---
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.
 |