CoolFace
Modelpublic

AiArtLab/sdxl_vae

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
17likes365downloads
Model Card

SDXL-VAE finetuned

=== Eval ===

SDXL VAE fp16 fix          | MSE=2.018e-03 PSNR=29.67 LPIPS=0.124 Edge=0.188 KL=32.222 | Z[min/mean/max/std]=[-4.066, -0.014, 4.301, 0.861] | Skew[min/mean/max]=[-0.017, 0.105, 0.165] | Kurt[min/mean/max]=[-0.380, -0.228, -0.108]
aiartlab/SDXLVAE           | MSE=1.736e-03 PSNR=30.29 LPIPS=0.116 Edge=0.181 KL=32.222 | Z[min/mean/max/std]=[-4.066, -0.014, 4.301, 0.861] | Skew[min/mean/max]=[-0.017, 0.105, 0.165] | Kurt[min/mean/max]=[-0.380, -0.228, -0.108]

=== Percent ===

| Модель                     |       MSE |      PSNR |     LPIPS |      Edge |
|----------------------------|-----------|-----------|-----------|-----------|
| SDXL VAE fp16 fix          |      100% |      100% |      100% |      100% |
| aiartlab/SDXLVAE           |    116.3% |    102.1% |    107.3% |    103.7% |

![Click it](https://imgsli.com/NDE1MjY1/1/2)

[image]

Diffusers

from diffusers import AutoencoderKL

vae = AutoencoderKL.from_pretrained("AiArtLab/sdxl_vae",subfolder="vae").cuda().half()

Train status, in progress:

We are currently testing the possibility of improving the SDXL VAE decoder by increasing its depth (asymmetric VAE). This will lead to a slight increase in model size (approximately 20 percent), but we expect this will improve reconstruction quality without modifying the encoder (does not require retraining SDXL). Unfortunately, our resources are quite limited (we train models on consumer GPUs, currently training three models: SDXL VAE, Simple Diffusion, and Simple VAE), so please be patient. Model training is a meticulous and time-consuming process.

VAE Training Process

  • —Encoder: Frozen (to avoid retraining SDXL for the new VAE).
  • —Dataset: 100,000 PNG images
  • —Training Time: 4 days
  • —Hardware: Single RTX 4090
  • —Resolution: 512px
  • —Precision: FP32
  • —Effective Batch Size: 16 (batch size 2 + gradient accumulation 8)
  • —Optimizer: AdamW (8-bit)
  • —++ MSE && Edge Loss: https://wandb.ai/recoilme/vae/runs/qy438uak

Implementation

  • —Base Code: Used a simple diffusion model training script.
  • —Training Target: Only the decoder, focusing on image reconstruction.

Loss Functions

  • —Initially used LPIPS and MSE.
  • —Noticed FID score improving, but images becoming blurry (FID overfits to blurry images—improving FID is not always good).
  • —Switched to MAE.
  • —Balanced LPIPS and MAE at 90/10 ratio.
  • —Used median perceptuallossweight for better balance.

Compare

https://imgsli.com/NDE1MjY1/1/2

Donations

Please contact with us if you may provide some GPU's or money on training

DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83

BTC: 3JHv9Hb8kEW8zMAccdgCdZGfrHeMhH1rpN

Contacts

recoilme