TAEF2

TAEF2

用一个微型自动编码器实时预览 FLUX.2 的图像生成，它与 FLUX.2 VAE 使用相同的潜空间格式。适合图像应用开发者和内存有限的人。

由 madebyollin 在 Hugging Face 上制作。这些说明是 Mindket 根据它的模型页面和模型卡写的，模型自己的文件和许可证在 Hugging Face 上。

简介

它是什么：TAEF2 是 madebyollin 做的 FLUX.2 微型自动编码器。它使用与 FLUX.2 VAE 相同的潜空间格式，所以可以替代它。模型卡点出两种用途：FLUX.2 生成时的实时预览，以及资源紧张时的编码或解码。这个仓库存放的是 safetensors 格式的权重。

适合谁：适合在 FLUX.2 上构建图像生成工具的开发者，以及在有限硬件上运行它的人。

行业 | 谁在用 | 怎么用 | 例子
软件 | 图像应用开发者 | 在 FLUX.2 生成时显示实时预览 | 在渲染完成前先预览
爱好 | 小显卡用户 | 用更少的内存编码和解码潜变量 | 用 CPU 卸载运行 FLUX.2 Klein
科研 | 模型研究者 | 对比微型和完整的自动编码器 | 把 TAEF2 和完整的 FLUX.2 VAE 对比

只想用现成应用生成图片的人可以跳过它，因为这是给流水线用的组件。

用在哪里：

• 装有 diffusers 库的 Python，通过模型卡给的包装代码
• FLUX.2 Klein 流水线，它在其中替代 VAE
• Windows、macOS 或 Linux，装有 Python，用来下载文件

使用场景 | 适合谁 | 你会得到什么
实时预览 | 应用开发者 | 生成过程中快速解码出的图像
低内存解码 | 爱好者 | 用小得多的模型解码出的图像
对比测试 | 研究者 | 与另外两个自动编码器的并排结果

如何设置：你需要先准备几样东西。

• 装有 torch 和 safetensors 的 Python
• 支持 FLUX.2 Klein 的 diffusers 版本，从源码安装
• taesd 项目里的 taesd.py 文件，包装代码会导入它
• 这个仓库里的 taef2.safetensors 权重

1. 安装支持 Klein 的 diffusers 版本。
2. 把 taesd.py 放在你的脚本旁边。
3. 下载 taef2.safetensors 权重。
4. 把模型卡里的包装代码粘贴到你的脚本里。

pip install -U huggingface_hub
hf download madebyollin/taef2

如何使用：TAEF2 还没有被正确集成进 diffusers，所以你要用一个小类把它包起来，让它看起来像 diffusers 的 VAE。然后把它赋给流水线的 vae，照常生成。

# 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

工具 | 作用
taesd.py | 提供包装代码所依赖的 TAESD 类
taef2.safetensors | 存放 TAEF2 权重
DiffusersTAEF2Wrapper | 让 TAEF2 看起来像 diffusers 的 VAE
Flux2KleinPipeline | FLUX.2 Klein 流水线，你要替换它的 vae
enable_sequential_cpu_offload | 把模型的一部分移到 CPU 以节省 GPU 内存

示例：你让流水线生成一块纽约风格的浆果芝士蛋糕，尺寸 1024 乘 1024，用 4 步。你换进去的解码器生成图片，脚本把它存为 flux-klein.png。

新手与专家：自动编码器在图片和图像模型使用的紧凑潜空间形式之间做转换。如果你是新手，按原样运行模型卡的示例，再看看 flux-klein.png。

因为 TAEF2 保持与 FLUX.2 VAE 相同的潜空间接口，你可以把它换进流水线而不用改别的东西。包装代码在转换 diffusers 权重时会移动解码器的层索引，使用 32 个潜通道和 flux_2 架构变体，并以 BF16 运行。

FLUX.2 流水线产生潜变量
包装器把它们交给 TAEF2 解码器
解码器返回缩放到负 1 到 1 的像素
图片被显示或保存为文件

注意事项：

• License: MIT。
• Cost to run: 没有提到 API 密钥或付费服务。示例会从 Hugging Face 下载 FLUX.2 Klein 4B，并使用 CUDA 设备。
• Data it reaches: 你的流水线产生的潜变量和图片。模型卡没有提到账号或网络访问。
• Not for: 不用包装代码就在 diffusers 里使用，因为模型卡说它还没有被正确集成。

Hugging Face 上的数据（2026年11月13日）

• 最近 30 天下载：0
• 点赞：32
• 许可证：mit
• Quantized versions: 1

链接

• 模型：https://huggingface.co/madebyollin/taef2
• 文件：https://huggingface.co/madebyollin/taef2/tree/main

获取模型

用 Hugging Face 的命令行下载它的文件：

pip install -U huggingface_hub
hf download madebyollin/taef2

然后按它的模型卡使用它。这个工具包里没有会自己运行的东西。