Add model and model card

README.md

@@ -0,0 +1,94 @@
+---
+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, 8, 64, 64)
+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.

config.json

@@ -0,0 +1,37 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.30.3",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": false,
+  "in_channels": 3,
+  "latent_channels": 8,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 32,
+  "scaling_factor": 0.9615,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d72453572b449c46d15a257872d65cecefb2fab05fdcd7d33fdcdb37f3ef8969
+size 334865516