TAE FLUX.2 · AcademiaSD

Tiny AutoEncoder (TAESD-style decoder) for fast, sharp live previews of FLUX.2 and FLUX.2 Klein in ComfyUI.

Real stock photos (Pixabay), 512×512 · Left: original · Center: TAE_Flux2_AcademiaSD · Right: Latent2RGB (ComfyUI's default preview for this model).
PSNR against the original photo, so it also includes what the VAE itself loses.

ComfyUI previews FLUX.2 with Latent2RGB, a linear 32→3 color projection that looks blurry and blocky. This decoder is distilled from the FLUX.2 VAE: it turns the sampler's latents into a real preview image, with no need for the full VAE.


File	`TAE_Flux2_AcademiaSD.safetensors`
Size	3.3 MB (fp16)
Parameters	1.66 M
Input	FLUX.2 latents, 128 channels at 1/16 resolution, as the sampler sees them (`x0`)
Output	RGB at 16× the latent resolution, range [0, 1]
PSNR	24.4 dB vs 16.4 dB for Latent2RGB (62 stock photos, 512×512, against the original photo)
Works with	FLUX.2 [dev] and FLUX.2 [klein] (4B and 9B): they share the same VAE

🚀 Usage in ComfyUI

Download TAE_Flux2_AcademiaSD.safetensors to ComfyUI/models/vae_approx/.
Install ComfyUI-KJNodes.
Add the Model Preview Override node between your FLUX.2 or Klein model and the sampler.
In its tiny_vae input, choose TAE_Flux2_AcademiaSD.safetensors.

ComfyUI's built-in TAESD preview method does not pick up this file: for FLUX.2's latent format it looks for a different file, taef2_decoder. Use the KJNodes node.

🔧 Training details

Teacher: FLUX.2 VAE (flux2-vae.safetensors).
Architecture: flat TAESD decoder, Clamp → conv → 4 × (3 blocks + 2× upsample + conv) → block → conv, width 64, 16× upscale.
Latent space: ComfyUI's Flux2 latent format (128 channels, already the sampler's x0 space), so no extra scaling is needed.
Data: varied images, 512×512 crops, encoded once with the real VAE.
Training: 30,000 steps, batch 8, 512×512 tiles, AdamW LR 5e-4 with cosine decay, L1 + FFT loss, EMA 0.999, latent noise augmentation for stable previews on the noisy x0 of the first steps.

⚠️ Limitations

Preview only: use the real VAE for the final decode.
Decoder only, there is no encoder.
Trained and tested with the FLUX.2 VAE only.

🇪🇸 Español

Tiny AutoEncoder (decoder tipo TAESD) para previsualizar FLUX.2 y FLUX.2 Klein en ComfyUI con nitidez y en tiempo real. Sustituye a Latent2RGB, la previsualización por defecto, borrosa y pixelada. Destilado del VAE de FLUX.2, que es el mismo en FLUX.2 [dev] y en Klein 4B y 9B: 24,4 dB de PSNR frente a 16,4 dB de Latent2RGB.

Uso:

Descarga TAE_Flux2_AcademiaSD.safetensors en ComfyUI/models/vae_approx/.
Instala ComfyUI-KJNodes.
Pon el nodo Model Preview Override entre el modelo de FLUX.2 o Klein y el sampler.
En su entrada tiny_vae, elige TAE_Flux2_AcademiaSD.safetensors.

El método de previsualización TAESD que trae ComfyUI no lo usa, porque para el formato de latente de FLUX.2 busca otro archivo (taef2_decoder). Usa el nodo de KJNodes.

Solo sirve para previsualizar: la imagen final se decodifica con el VAE real.

🙏 Credits

TAESD architecture: madebyollin/taesd
Model Preview Override node: kijai/ComfyUI-KJNodes
FLUX.2: Black Forest Labs

🔗 AcademiaSD

YouTube · X / Twitter · Discord · Ko-fi

Downloads last month: -; Downloads are not tracked for this model. How to track