TAE Krea 2 · AcademiaSD
Tiny AutoEncoder (TAESD-style decoder) for fast, sharp live previews of Krea 2 in ComfyUI.
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:
1. Descarga TAE_Krea2_AcademiaSD.safetensors en ComfyUI/models/vae_approx/.
2. Instala ComfyUI-KJNodes.
3. Pon el nodo Model Preview Override entre el modelo de Krea 2 y el sampler.
4. 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
AcademiaSD
YouTube · X / Twitter · Discord · Ko-fi
