TAE Krea 2 · AcademiaSD

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

Example

Real stock photos (Pixabay), 512×512 · Left: original · Center: TAE_Krea2_AcademiaSD · Right: Latent2RGB (ComfyUI's default preview for this model).
PSNR against the original photo, so it also includes what the VAE itself loses: these detailed photos score lower than the validation figure below.

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


File	`TAE_Krea2_AcademiaSD.safetensors`
Size	2.5 MB (fp16)
Parameters	1.23 M
Input	Krea 2 latents, 16 channels, as the sampler sees them (`x0`)
Output	RGB at 8× the latent resolution, range [0, 1]
Decode time	~12 ms for a 1024×1024 preview (RTX 5080, fp16)
Validation PSNR	36.1 dB vs 22.8 dB for Latent2RGB

🚀 Usage in ComfyUI

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

ComfyUI's built-in TAESD preview method does not pick up this file: for Krea 2's latent format (Wan21) it looks for a different file, the lighttaew2_1 video decoder. Use the KJNodes node.

🔧 Training details

Teacher: Krea 2 VAE (krea2RealVae_v10.safetensors).
Architecture: flat TAESD decoder, Clamp → conv → 3 × (3 blocks + 2× upsample + conv) → block → conv, width 64.
Latent space: latents normalized with ComfyUI's Wan21 latent format (mean/std), the one ComfyUI uses for Krea 2 and the same space as the sampler's x0, so no extra scaling is needed.
Data: 2,048 varied images, 512×512 crops, encoded once with the real VAE.
Training: 30,000 steps, batch 8, 256×256 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 Krea 2 only.

🇪🇸 Español

Tiny AutoEncoder (decoder tipo TAESD) para previsualizar Krea 2 en ComfyUI con nitidez y en tiempo real. Sustituye a Latent2RGB, la previsualización por defecto, borrosa y pixelada. Destilado del VAE de Krea 2: decodifica una previsualización de 1024 px en unos 12 ms, con 36,1 dB de PSNR frente a 22,8 dB de Latent2RGB.

Uso:

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

El método de previsualización TAESD que trae ComfyUI no lo usa, porque para el formato de latente de Krea 2 busca otro archivo. 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
Krea 2: Krea team

🔗 AcademiaSD

YouTube · X / Twitter · Discord · Ko-fi

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