Hugging Face's logo

SD-Next
/
MicroDecoder

Text-to-Image
Diffusers
vae
Model card Files
xet
 Community
MicroDecoder
File size: 4,806 Bytes
f9898d3
 
 
 
f7adadd
 
f9898d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3991418
 
 
 
978e897
 
 
3991418
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
---
license: apache-2.0
pipeline_tag: text-to-image
library_name: diffusers
tags:
- vae
---

# 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](https://github.com/vladmandic/sdnext/blob/dev/modules/vae/sd_vae_micro_model.py)
- Training code [here](https://github.com/vladmandic/sdnext/blob/dev/modules/vae/sd_vae_micro_train.py)
- Example inference code [here](https://github.com/vladmandic/sdnext/blob/dev/modules/vae/sd_vae_micro.py)

## Example

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

![MicroDecoder](https://huggingface.co/vladmandic/MicroDecoder/resolve/main/MicroDecoder.jpg)

## Training

```shell
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
```

```log
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?

```json
    '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'],
```