TAE FLUX.2 - IDEOGRAM-4.0 · AcademiaSD
Tiny AutoEncoder (TAESD-style decoder) for fast, sharp live previews of FLUX.2, FLUX.2 Klein and Ideogram 4.0 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.safetensorstoComfyUI/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_vaeinput, chooseTAE_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
Flux2latent format (128 channels, already the sampler'sx0space), 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
x0of 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.safetensorsenComfyUI/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, eligeTAE_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