Update README.md

README.md

@@ -1,3 +1,38 @@
----
-license: mit
----
+---
+license: mit
+---
+
+# Transfusion - VAE
+
+## How to use with 🧨 diffusers
+
+```py
+from diffusers.models import AutoencoderKL
+vae = AutoencoderKL.from_pretrained("lavinal712/transfusion-vae")
+```
+
+## Model 
+
+This model was trained for 7 epochs on ImageNet, with training parameters following the original Transfusion paper.
+
+$\mathcal{L}_{\mathrm{VAE}} = \mathcal{L}_1 + \mathcal{L}_{\mathrm{LPIPS}} + 0.5\mathcal{L}_{\mathrm{GAN}} + 0.2\mathcal{L}_{\mathrm{ID}} + 0.000001\mathcal{L}_{\mathrm{KL}}$
+
+## Evaluation 
+
+ImageNet 2012 (256x256, val, 50000 images)
+
+| Model           | rFID  | PSNR   | SSIM  | LPIPS |
+|-----------------|-------|--------|-------|-------|
+| Transfusion-VAE | 0.567 | 28.195 | 0.829 | 0.100 |
+| SD-VAE          | 0.692 | 26.910 | 0.772 | 0.130 |
+
+PSNR: 28.19461106581366 ± 4.319437641910872
+SSIM: 0.8292981386184692 ± 0.10117273032665253
+LPIPS: 0.10047737887922209 ± 0.04119070588416227
+rFID: 0.5673381146947349
+
+Paper: [Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model](https://arxiv.org/abs/2408.11039)
+
+Dataset: [ImageNet](https://image-net.org/)
+
+Base Code: [lavinal712/AutoencoderKL](https://github.com/lavinal712/AutoencoderKL)