README.md

---
license: cc0-1.0
datasets:
- nyanko7/danbooru2023
- boxingscorpionbagel/e621-2024
library_name: diffusers
---
# LibreVAE

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
from diffusers import AutoencoderKL
from PIL import Image
from torchvision import transforms

transform_image = transforms.Compose([
    transforms.Lambda(lambda x: x.convert("RGB")),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

untransform_image = transforms.Compose([
    transforms.Normalize((-1, -1, -1), (2, 2, 2)),
    transforms.ToPILImage()
])

model = AutoencoderKL.from_pretrained("scrumptious/librevae-f8-d8").to("cuda").eval()
model.requires_grad_(False)

# For a 512x512 image, image_tensor will be (1, 3, 512, 512)
image_tensor = transform_image(Image.open("sample.png")).unsqueeze(0).to("cuda")
# latent will be (1, 8, 64, 64)
# we multiply it by the scaling factor so it has an approximate mean of 0 and variance of 1
latent = model.encode(image_tensor).latent_dist.sample() * model.config.scaling_factor
# output will be (1, 3, 512, 512)
output = model.decode(latent / model.config.scaling_factor).sample
output_image = untransform_image(output.squeeze(0).clamp(-1, 1).cpu())

output_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](https://github.com/NovelAI/novelai-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.

### Other Details

- **Precision:** BF16 mixed precision
- **Learning Rate:** 1e-4 with a 50% decay per epoch
- **Epochs:** 2
- **Optimizer:** AdamW
- **Batch Size:** 2 (per-GPU batch) \* 2 (GPUs) \* 128 (gradient accumulation steps) = 512

### Validation Performance
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](https://creativecommons.org/publicdomain/zero/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.