MicroDecoder

MicroDecoder VAE can be trained on any diffusion model in ~5-15min and provides low-quality reconstruction from latents to RGB in ~0.01sec.

Inference includes noise correction based on current timestep, blurring corrections 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.

Code

• Model definition here
• Training code here
• Example inference code here

Example

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

Training

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

MicroDecoder Train
Args: Namespace(folder='/home/vlado/generative/Input/vae/', max=500, resolution=512, crop='center', vae='AutoencoderKLMiniMaxH3', repo='OzzyGT/MiniMax_H3_sdnq_dynamic_4bit', subfolder='vae', dim=256, epochs=400, scale=4, batch=16, lr=0.0003, ema=0.999, val=0.1, output='microdecoder-minimaxh3.safetensors', device='cuda')
Base VAE: cls=<class 'diffusers.models.autoencoders.autoencoder_kl_minimax_h3.AutoencoderKLMiniMaxH3'> repo="OzzyGT/MiniMax_H3_sdnq_dynamic_4bit" subfolder="vae"
Init VAE: model=MicroDecoder ema=EMAModel optimizer=AdamW scheduler=CosineAnnealingLR criterion=EnhancedQualityLoss metrics=DetailedMetricsTracker
Epoch   | tTotal    | tPSNR   | tSSIM   | vTotal   | vPSNR    | vSSIM    | vL1     | vLAB
-------------------------------------------------------------------------------------------
001/400 | 7.0891    | 10.49   | 0.2540  | 7.1642   | 10.75    | 0.2357   | 0.2614  | 0.4146
...
400/400 | 4.5365    | 21.03   | 0.6402  | 1.9644   | 30.61    | 0.9266   | 0.0204  | 0.0465
Load   ━━━━━━━━━━━━━━━━━━━━ 400/400 100% 0:00:00 0:00:08 Images=400 Shape=torch.Size([400, 3, 512, 512])
Encode ━━━━━━━━━━━━━━━━━━━━ 400/400 100% 0:00:00 0:00:46 Train=360 Validation=40 EMA=0.999
Train  ━━━━━━━━━━━━━━━━━━━━ 400/400 100% 0:00:00 0:06:33 Train(psnr=21.029 ssim=0.640 lpips=2.784 l1=0.089 lab=0.156 grad=0.073 sat=0.124 fft=0.071 total=4.536) Validate(psnr=30.607 ssim=0.927 lpips=1.242 l1=0.020 lab=0.046 grad=0.046 sat=0.057 fft=0.032 total=1.964)
Complete: Time=393.89 Epoch/Sec=1.02 PSNR=30.61 dB @ Epoch=400
Save: filename="microdecoder-minimaxh3.safetensors"

Models

MicroDecoder includes following pre-trained models:

• sd, sdxl, f1, f2, sd3, wan21, wan22, minimaxh3, qwen21, dcae, mage, cog, hy

And each VAE model can be used for different base models as long as they share the same VAE architecture

VAE to Model Map

Which MicroDecoder model can be used for which model types?

    'sd': ['sd',
           'sdxs'],
    'sdxl': ['sdxl',
             'ldm',
             'pixartalpha',
             'pixartsigma',
             'hunyuandit',
             'omnigen',
             'auraflow',
             'bria',
             'kolors',
             'lumina'],
    'f1': ['f1',
           'h1',
           'zimage',
           'lumina2',
           'chroma',
           'longcat',
           'omnigen2',
           'flite',
           'ovis',
           'kandinsky5',
           'ultraflux',
           'boogu',
           'step1x_edit'],
    'f2': ['f2',
           'f3',
           'ernieimage',
           'lens',
           'ideogram4',
           'lladaimage',
           'mossimage',
           'sefi',
           'fibo2'],
    'sd3': ['sd3'],
    'wan21': ['wanai',
              'qwen',
              'chrono',
              'cosmos',
              'anima',
              'joy',
              'krea2',
              'mingimage',
              'nucleusimage'],
    'wan22': ['wan22',
              'fibo'],
    'minimaxh3': ['minimaxh3'],
    'qwen21': ['qwen21'],
    'dcae': ['sana'],
    'mage': ['mageflow'],
    'cog': ['cogview3',
            'cogview4'],
    'hy': ['hunyuanimage'],