MicroDecoder

Micro-VAE that can be trained on any diffusion model in ~15min and provides low-quality reconstruction from latents to RGB in ~0.01sec.
Inference includes noise correction based on current timestep, intentional blurring plus upscale interpolation: all with intention of providing as fast-as-possible reconstruction that is viable and consistent with any noise levels.
Intended use-case is live-preview during generative model inference.

Shapes/Channels/etc are inferred from the base VAE, so no configuration changes are needed between different models.

Model definition and training code here
Example inference code here

Example using MicroDecoder with Flux.2-Klein-9B and compared with official final VAE processing at the end:

Example

sd_vae_micro_train.py \
  --dim 256 \
  --epochs 350 \
  --resolution 512 \
  --scale 4 \
  --lr 0.0003
  --folder ~/generative/Input/vae/ \
  --vae AutoencoderKLQwenImage21 \
  --repo Qwen/Qwen-Image-2.1
  --output MicroVAE-qwen21.safetensors

Encoding Latents: 100%
Training MicroDecoder [350 epochs | Channels: 32 | Hidden: 256 | Scale: 4x]
Dataset: 360 train samples | 40 validation samples | EMA decay: 0.999

Epoch  | Train Tot | Tr PSNR | Tr SSIM | Val Tot  | Val PSNR | Val SSIM | Val L1  | Val LAB
--------------------------------------------------------------------------------------------
001/350  | 5.7244    | 10.38   | 0.3431  | 5.9729   | 10.56    | 0.3047   | 0.2676  | 0.4254  *
100/350  | 2.7690    | 25.49   | 0.8030  | 2.2011   | 26.52    | 0.8672   | 0.0376  | 0.0699  *
200/350  | 2.4006    | 27.04   | 0.8400  | 1.1679   | 34.24    | 0.9635   | 0.0132  | 0.0317  *
300/350  | 2.3950    | 28.05   | 0.8472  | 1.0002   | 35.63    | 0.9723   | 0.0113  | 0.0282  *
350/350  | 2.3010    | 28.51   | 0.8511  | 0.9781   | 35.83    | 0.9734   | 0.0111  | 0.0278  *

Training complete! Best validation PSNR: 35.83 dB (Epoch 350)
Saved best EMA model weights to: MicroVAE-qwen21.safetensors

Downloads last month: 2