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

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 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 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.

scrumptious
/

librevae-f8-d8