LibreVAE is a Variational Autoencoder designed to serve as a component for future generative modelling projects. It has 8 latent channels, and reduces images dimensionally by a factor of 8. It was trained using HuggingFace Diffusers, and can be loaded with the AutoencoderKL class.
Example Usage
python
1from diffusers import AutoencoderKL
2from PIL import Image
3from torchvision import transforms
45transform_image = transforms.Compose([6 transforms.Lambda(lambda x: x.convert("RGB")),7 transforms.ToTensor(),8 transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))9])1011untransform_image = transforms.Compose([12 transforms.Normalize((-1,-1,-1),(2,2,2)),13 transforms.ToPILImage()14])1516model = AutoencoderKL.from_pretrained("scrumptious/librevae-f8-d8").to("cuda").eval()17model.requires_grad_(False)1819# For a 512x512 image, image_tensor will be (1, 3, 512, 512)20image_tensor = transform_image(Image.open("sample.png")).unsqueeze(0).to("cuda")21# latent will be (1, 8, 64, 64)22# we multiply it by the scaling factor so it has an approximate mean of 0 and variance of 123latent = model.encode(image_tensor).latent_dist.sample()* model.config.scaling_factor
24# output will be (1, 3, 512, 512)25output = model.decode(latent / model.config.scaling_factor).sample
26output_image = untransform_image(output.squeeze(0).clamp(-1,1).cpu())2728output_image.save('sample_decoded.png')
Training Details
Training Datasets
LibreVAE was trained on the e621-2024 and danbooru2023 datasets, both of which are large, curated collections of artwork. While the model was trained primarily on artwork, our testing showed that it was capable of working with other types of images.
Dataset Preprocessing
We applied a modified version of NovelAI's Aspect Ratio Bucketing to the images, where we dynamically selected aspect ratio buckets using K-Means over our training dataset instead of predetermining them based on fixed sizes. We then set sizes for these buckets to be around 256x256.
Loss Function
The loss function for this model was MSE_lab + (0.5 * MSE_rgb) + (0.1 * LPIPS) + (1e-4 * KL), where MSE_lab was the mean squared error calculated in CIELAB color space, MSE_rgb was the mean squared error calculated in RGB color space, LPIPS was the LPIPS loss, and KL was the KL divergence.
The model achieved the following scores in its final validation run.
MSE in CIELAB space: 0.002154
MSE in RGB space: 0.0062
LPIPS: 0.0555
Uses
LibreVAE is intended to be used by researchers or developers as a component for generative models, such as text-to-image models. The developers don't forsee any direct uses that would not be better served by an existing image compression solution.
License
The weights for LibreVAE are released under the CC0 1.0 license.
Citation
Under the CC0 1.0 license, you are not required to provide any attribution when using or redistributing LibreVAE. If you use LibreVAE in your research or projects and would like to provide attribution, you can cite it as:
@misc{LibreVAE2024,
title={LibreVAE},
author={Scrumptious AI Labs},
year={2024},
note={https://huggingface.co/scrumptious/librevae-f8-d8}
}
Acknowledgments
Special thanks to the contributors of the e621-2024 and danbooru2023 datasets, the HuggingFace Diffusers team, and the PyTorch Lightning team.
