|
--- |
|
license: mit |
|
pipeline_tag: image-to-image |
|
--- |
|
|
|
<h2 align="center"> <a href="https://arxiv.org/abs/2503.14325">LeanVAE: An Ultra-Efficient Reconstruction VAE for Video Diffusion Models</a></h2> |
|
|
|
https://github.com/user-attachments/assets/a2a4814a-192b-4cc4-b1a3-d612caa1d872 |
|
|
|
We present **LeanVAE**, a lightweight Video VAE designed for ultra-efficient video compression and scalable generation in Latent Video Diffusion Models (LVDMs). |
|
|
|
- **Lightweight & Efficient**: Only **40M parameters**, significantly reducing computational overhead 📉 |
|
- **Optimized for High-Resolution Videos**: Encodes and decodes a **17-frame 1080p video** in **3 seconds** using only **15GB of GPU memory** *(without tiling inference)* 🎯 |
|
- **State-of-the-Art Video Reconstruction**: Competes with leading Video VAEs 🏆 |
|
- **Versatile**: Supports both **images and videos**, preserving **causality in latent space** 📽️ |
|
- **Evidenced by Diffusion Model**: Enhances visual quality in video generation ✨ |
|
|
|
--- |
|
## 🛠️ **Installation** |
|
Clone the repository and install dependencies: |
|
``` |
|
git clone https://github.com/westlake-repl/LeanVAE |
|
cd LeanVAE |
|
pip install -r requirements.txt |
|
``` |
|
--- |
|
## 🎯 **Quick Start** |
|
**Train LeanVAE** |
|
```bash |
|
bash scripts/train.sh |
|
``` |
|
|
|
**Run Video Reconstruction** |
|
```bash |
|
bash scripts/inference.sh |
|
``` |
|
|
|
**Evaluate Reconstruction Quality** |
|
```bash |
|
bash scripts/eval.sh |
|
``` |
|
--- |
|
|
|
## 📜 **Pretrained Models** |
|
### Video VAE Model: |
|
| Model | PSNR ⬆️ | LPIPS ⬇️ | Params 📦 | TFLOPs ⚡ | Checkpoint 📥 | |
|
| ---------------- | ------ | ------- | -------- | -------- | ----------------------------------- | |
|
| **LeanVAE-4ch** | 26.04 | 0.0899 | 39.8M | 0.203 | [LeanVAE-chn4.ckpt](https://huggingface.co/Yumic/LeanVAE/resolve/main/LeanVAE-dim4.ckpt?download=true) | |
|
| **LeanVAE-16ch** | 30.15 | 0.0461 | 39.8M | 0.203 | [LeanVAE-chn16.ckpt](https://huggingface.co/Yumic/LeanVAE/resolve/main/LeanVAE-dim16.ckpt?download=true) | |
|
|
|
|
|
### Latte Model: |
|
The code and pretrained weights for video generation will be released soon. Stay tuned! |
|
| Model | Dataset | FVD ⬇️ | Checkpoint 📥 | |
|
| ---------- | ---------- | ---------- | ----------- | |
|
| Latte + LeanVAE-chn4 | SkyTimelapse |49.59 | sky-chn4.ckpt | |
|
| Latte + LeanVAE-chn4 | UCF101 |164.45 | ucf-chn4.ckpt | |
|
| Latte + LeanVAE-chn16 | SkyTimelapse |95.15 | sky-chn16.ckpt | |
|
| Latte + LeanVAE-chn16 | UCF101 |175.33 | ucf-chn16.ckpt | |
|
|
|
--- |
|
## 🔧 **Using LeanVAE in Your Project** |
|
|
|
```python |
|
from LeanVAE import LeanVAE |
|
|
|
# Load pretrained model |
|
model = LeanVAE.load_from_checkpoint("path/to/ckpt", strict=False) |
|
|
|
# 🔄 Encode & Decode an Image |
|
image, image_rec = model.inference(image) |
|
|
|
# 🖼️ Encode an image → Get latent : |
|
latent = model.encode(image) # (B, C, H, W) → (B, d, 1, H/8, W/8), where d=4 or 16 |
|
|
|
# 🖼️ Decode latent representation → Reconstruct image |
|
image = model.decode(latent, is_image=True) # (B, d, 1, H/8, W/8) → (B, C, H, W) |
|
|
|
|
|
# 🔄 Encode & Decode a Video |
|
video, video_rec = model.inference(video) ## Frame count must be 4n+1 (e.g., 5, 9, 13, 17...) |
|
|
|
# 🎞️ Encode Video → Get Latent Space |
|
latent = model.encode(video) # (B, C, T+1, H, W) → (B, d, T/4+1, H/8, W/8), where d=4 or 16 |
|
|
|
# 🎞️ Decode Latent → Reconstruct Video |
|
video = model.decode(latent) # (B, d, T/4+1, H/8, W/8) → (B, C, T+1, H, W) |
|
|
|
# ⚡ Enable **Temporal Tiling Inference** for Long Videos |
|
model.set_tile_inference(True) |
|
model.chunksize_enc = 5 |
|
model.chunksize_dec = 5 |
|
``` |
|
--- |
|
|
|
## 📂 **Preparing Data for Training** |
|
To train LeanVAE, you need to create metadata files listing the video paths, grouped by resolution. Each file contains paths to videos of the same resolution. |
|
``` |
|
📂 data_list |
|
├── 📄 96x128.txt 📜 # Contains paths to all 96x128 videos |
|
│ ├── /path/to/video_1.mp4 |
|
│ ├── /path/to/video_2.mp4 |
|
│ ├── ... |
|
├── 📄 256x256.txt 📜 # Contains paths to all 256×256 videos |
|
│ ├── /path/to/video_3.mp4 |
|
│ ├── /path/to/video_4.mp4 |
|
│ ├── ... |
|
├── 📄 352x288.txt 📜 # Contains paths to all 352x288 videos |
|
│ ├── /path/to/video_5.mp4 |
|
│ ├── /path/to/video_6.mp4 |
|
│ ├── ... |
|
``` |
|
📌 Each text file lists video paths corresponding to a specific resolution. Set `args.train_datalist` to the folder containing these files. |
|
|
|
|
|
--- |
|
## 📜 **License** |
|
|
|
This project is released under the **MIT License**. See the `LICENSE` file for details. |
|
|
|
|
|
## 🔥 **Why Choose LeanVAE?** |
|
LeanVAE is **fast, lightweight and powerful**, enabling high-quality video compression and generation with minimal computational cost. |
|
|
|
If you find this work useful, consider **starring ⭐ the repository** and citing our paper! |
|
|
|
--- |
|
|
|
## 📝 **Cite Us** |
|
```bibtex |
|
@misc{cheng2025leanvaeultraefficientreconstructionvae, |
|
title={LeanVAE: An Ultra-Efficient Reconstruction VAE for Video Diffusion Models}, |
|
author={Yu Cheng and Fajie Yuan}, |
|
year={2025}, |
|
eprint={2503.14325}, |
|
archivePrefix={arXiv}, |
|
primaryClass={cs.CV}, |
|
url={https://arxiv.org/abs/2503.14325}, |
|
} |
|
``` |
|
--- |
|
|
|
## 👍 **Acknowledgement** |
|
Our work benefits from the contributions of several open-source projects, including [OmniTokenizer](https://github.com/FoundationVision/OmniTokenizer), [Open-Sora-Plan](https://github.com/PKU-YuanGroup/Open-Sora-Plan), [VidTok](https://github.com/microsoft/VidTok), and [Latte](https://github.com/Vchitect/Latte). We sincerely appreciate their efforts in advancing research and open-source collaboration! |
