newvae

README.md

@@ -85,7 +85,27 @@ image.show()
 Prompt: *cat* (Neg: low quality, bad quality, photo, low details)
 ![cat](media/cat.jpg)
 
-### Image upscale
+### VAE
+
+The VAE in Simple Diffusion utilizes an asymmetric VAE architecture featuring an 8x encoder and a 16x decoder. While a compression factor of 8 is maintained during training, the resolution is effectively doubled during inference through an additional upscaling block. This strategy reduces training costs by an order of magnitude and boosts inference speed without perceptual quality loss. Effectively, this acts as an integrated latent upscaler. To ensure a fair comparison with other VAEs, we downsampled the generated images to match the input resolution for metric evaluation. The SDXS VAE was not trained from scratch but was initialized from weights of FLUX 2 VAE, then redisigned and retrained. We also trained [16 ch vae](https://huggingface.co/AiArtLab/simplevae) with flux.1 quality based on aura vae.
+
+```
+Wan2.2-T2V-A14B            | MSE=8.908e-04 PSNR=33.21 LPIPS=0.052
+FLUX.1-schnell VAE         | MSE=6.020e-04 PSNR=34.59 LPIPS=0.038
+FLUX.2-klein-9B            | MSE=3.638e-04 PSNR=36.60 LPIPS=0.025
+sdxs-1b                    | MSE=3.530e-04 PSNR=36.92 LPIPS=0.027
+
+=== Percents ===
+| Модель                     |       MSE |      PSNR |     LPIPS |      Edge |    Skew|0 |    Kurt|0 |
+|----------------------------|-----------|-----------|-----------|-----------|-----------|-----------|
+| Wan2.2-T2V-A14B            |      100% |      100% |      100% |      100% |      100% |      100% |
+| FLUX.1-schnell VAE         |    148.0% |    104.2% |    136.7% |    128.8% |    104.9% |    112.2% |
+| FLUX.2-klein-9B            |    244.8% |    110.2% |    209.7% |    165.9% |     93.2% |    118.2% |
+| sdxs-1b                    |    252.4% |    111.1% |    193.1% |    172.1% |     93.2% |    118.3% |
+```
+To the best of our knowledge, this VAE achieves SOTA: [eval.py](src/eval.py)
+
+### Image upscale (vae)
 ```
 upscaled = pipe.image_upscale("media/girl.jpg")
 upscaled[0].show()
@@ -229,25 +249,6 @@ image = pipeline(
     )[0]
 ```
 
-### VAE
-
-The VAE in Simple Diffusion utilizes an asymmetric VAE architecture featuring an 8x encoder and a 16x decoder. While a compression factor of 8 is maintained during training, the resolution is effectively doubled during inference through an additional upscaling block. This strategy reduces training costs by an order of magnitude and boosts inference speed without perceptual quality loss. Effectively, this acts as an integrated latent upscaler. To ensure a fair comparison with other VAEs, we downsampled the generated images to match the input resolution for metric evaluation. The SDXS VAE was not trained from scratch but was initialized from weights of FLUX 2 VAE, then redisigned and retrained. We also trained [16 ch vae](https://huggingface.co/AiArtLab/simplevae) with flux.1 quality based on aura vae.
-
-```
-Wan2.2-T2V-A14B            | MSE=8.908e-04 PSNR=33.21 LPIPS=0.052
-FLUX.1-schnell VAE         | MSE=6.020e-04 PSNR=34.59 LPIPS=0.038
-FLUX.2-klein-9B            | MSE=3.638e-04 PSNR=36.60 LPIPS=0.025
-sdxs-1b                    | MSE=3.491e-04 PSNR=36.95 LPIPS=0.027
-
-=== Percents ===
-| Модель                     |       MSE |      PSNR |     LPIPS |      Edge |    Skew|0 |    Kurt|0 |
-|----------------------------|-----------|-----------|-----------|-----------|-----------|-----------|
-| Wan2.2-T2V-A14B            |      100% |      100% |      100% |      100% |      100% |      100% |
-| FLUX.1-schnell VAE         |    148.0% |    104.2% |    136.7% |    128.8% |    104.9% |    112.2% |
-| FLUX.2-klein-9B            |    244.8% |    110.2% |    209.7% |    165.9% |     93.2% |    118.3% |
-| sdxs-1b                    |    255.2% |    111.2% |    193.7% |    171.6% |     93.3% |    118.2% |
-```
-To the best of our knowledge, this VAE achieves SOTA: [eval.py](src/eval.py)
 
 ### Unet
 
@@ -271,8 +272,7 @@ Training was conducted with a maximum of 250 tokens, and a 10% dropout rate was
 Additionally, the use of a full-fledged language model allowed us to integrate an optional prompt enhancement mechanism into the pipeline.
 
 ### Retrospective and Key Takeaways
-## The Journey Begins
-This adventure started in December 2024 after the release of the SANA model. We received a donation from Stan for fine-tuning SANA and, together with Stas, began fine-tuning and further developing it. Despite spending the entire budget, we did not achieve significant improvements. However, we were shocked by how poorly the model was trained and designed, and we became convinced that we could do better—though we were wrong.
+Research started in December 2024 after the release of the SANA model. We received a donation from Stan for fine-tuning SANA and, together with Stas, began fine-tuning and further developing it. Despite spending the entire budget, we did not achieve significant improvements. However, we were shocked by how poorly the model was trained and designed, and we became convinced that we could do better—though we were wrong.
 Shifting Gears
 By February 2025, we split our efforts and began designing our own architectures—which we are still doing today. Stas favored the DiT architecture, while I believed in UNet. Despite some differences in architectural views, we maintained close communication, shared our work, and supported each other throughout the process. We also engaged with the AIArtLab community (a virtual Telegram chat for those contributing to model development)—thank you all for your support.
 ## Main mistake

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:14a8b19e30dd88dc0ecaaa8ee1fef716e6b0570b7e9133af42feb6465f6af25a
-size 17570487
+oid sha256:ffb8b63e9a963aee212f993c38154f7fef82b7c94325cb2266c39093d9aad108
+size 13792585

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:780993a422cf0d783088c5116063ab6e3b1453a5256f205b78b1c67617f17b33
+oid sha256:d0963803789980cdf3df4d4e87b26c099ccd3ed08b78deb7a12afaba5e193c97
 size 383499124