Add pipeline tag to model card
Browse filesThis PR adds the `pipeline_tag: image-to-image` to the model card metadata. This makes the model discoverable through the model hub interface.
README.md
CHANGED
@@ -1,3 +1,144 @@
|
|
1 |
-
---
|
2 |
-
license: mit
|
3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
license: mit
|
3 |
+
pipeline_tag: image-to-image
|
4 |
+
---
|
5 |
+
|
6 |
+
<h2 align="center"> <a href="https://arxiv.org/abs/2503.14325">LeanVAE: An Ultra-Efficient Reconstruction VAE for Video Diffusion Models</a></h2>
|
7 |
+
|
8 |
+
https://github.com/user-attachments/assets/a2a4814a-192b-4cc4-b1a3-d612caa1d872
|
9 |
+
|
10 |
+
We present **LeanVAE**, a lightweight Video VAE designed for ultra-efficient video compression and scalable generation in Latent Video Diffusion Models (LVDMs).
|
11 |
+
|
12 |
+
- **Lightweight & Efficient**: Only **40M parameters**, significantly reducing computational overhead π
|
13 |
+
- **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)* π―
|
14 |
+
- **State-of-the-Art Video Reconstruction**: Competes with leading Video VAEs π
|
15 |
+
- **Versatile**: Supports both **images and videos**, preserving **causality in latent space** π½οΈ
|
16 |
+
- **Evidenced by Diffusion Model**: Enhances visual quality in video generation β¨
|
17 |
+
|
18 |
+
---
|
19 |
+
## π οΈ **Installation**
|
20 |
+
Clone the repository and install dependencies:
|
21 |
+
```
|
22 |
+
git clone https://github.com/westlake-repl/LeanVAE
|
23 |
+
cd LeanVAE
|
24 |
+
pip install -r requirements.txt
|
25 |
+
```
|
26 |
+
---
|
27 |
+
## π― **Quick Start**
|
28 |
+
**Train LeanVAE**
|
29 |
+
```bash
|
30 |
+
bash scripts/train.sh
|
31 |
+
```
|
32 |
+
|
33 |
+
**Run Video Reconstruction**
|
34 |
+
```bash
|
35 |
+
bash scripts/inference.sh
|
36 |
+
```
|
37 |
+
|
38 |
+
**Evaluate Reconstruction Quality**
|
39 |
+
```bash
|
40 |
+
bash scripts/eval.sh
|
41 |
+
```
|
42 |
+
---
|
43 |
+
|
44 |
+
## π **Pretrained Models**
|
45 |
+
### Video VAE Model:
|
46 |
+
| Model | PSNR β¬οΈ | LPIPS β¬οΈ | Params π¦ | TFLOPs β‘ | Checkpoint π₯ |
|
47 |
+
| ---------------- | ------ | ------- | -------- | -------- | ----------------------------------- |
|
48 |
+
| **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) |
|
49 |
+
| **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) |
|
50 |
+
|
51 |
+
|
52 |
+
### Latte Model:
|
53 |
+
The code and pretrained weights for video generation will be released soon. Stay tuned!
|
54 |
+
| Model | Dataset | FVD β¬οΈ | Checkpoint π₯ |
|
55 |
+
| ---------- | ---------- | ---------- | ----------- |
|
56 |
+
| Latte + LeanVAE-chn4 | SkyTimelapse |49.59 | sky-chn4.ckpt |
|
57 |
+
| Latte + LeanVAE-chn4 | UCF101 |164.45 | ucf-chn4.ckpt |
|
58 |
+
| Latte + LeanVAE-chn16 | SkyTimelapse |95.15 | sky-chn16.ckpt |
|
59 |
+
| Latte + LeanVAE-chn16 | UCF101 |175.33 | ucf-chn16.ckpt |
|
60 |
+
|
61 |
+
---
|
62 |
+
## π§ **Using LeanVAE in Your Project**
|
63 |
+
|
64 |
+
```python
|
65 |
+
from LeanVAE import LeanVAE
|
66 |
+
|
67 |
+
# Load pretrained model
|
68 |
+
model = LeanVAE.load_from_checkpoint("path/to/ckpt", strict=False)
|
69 |
+
|
70 |
+
# π Encode & Decode an Image
|
71 |
+
image, image_rec = model.inference(image)
|
72 |
+
|
73 |
+
# πΌοΈ Encode an image β Get latent :
|
74 |
+
latent = model.encode(image) # (B, C, H, W) β (B, d, 1, H/8, W/8), where d=4 or 16
|
75 |
+
|
76 |
+
# πΌοΈ Decode latent representation β Reconstruct image
|
77 |
+
image = model.decode(latent, is_image=True) # (B, d, 1, H/8, W/8) β (B, C, H, W)
|
78 |
+
|
79 |
+
|
80 |
+
# π Encode & Decode a Video
|
81 |
+
video, video_rec = model.inference(video) ## Frame count must be 4n+1 (e.g., 5, 9, 13, 17...)
|
82 |
+
|
83 |
+
# ποΈ Encode Video β Get Latent Space
|
84 |
+
latent = model.encode(video) # (B, C, T+1, H, W) β (B, d, T/4+1, H/8, W/8), where d=4 or 16
|
85 |
+
|
86 |
+
# ποΈ Decode Latent β Reconstruct Video
|
87 |
+
video = model.decode(latent) # (B, d, T/4+1, H/8, W/8) β (B, C, T+1, H, W)
|
88 |
+
|
89 |
+
# β‘ Enable **Temporal Tiling Inference** for Long Videos
|
90 |
+
model.set_tile_inference(True)
|
91 |
+
model.chunksize_enc = 5
|
92 |
+
model.chunksize_dec = 5
|
93 |
+
```
|
94 |
+
---
|
95 |
+
|
96 |
+
## π **Preparing Data for Training**
|
97 |
+
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.
|
98 |
+
```
|
99 |
+
π data_list
|
100 |
+
βββ π 96x128.txt π # Contains paths to all 96x128 videos
|
101 |
+
β βββ /path/to/video_1.mp4
|
102 |
+
β βββ /path/to/video_2.mp4
|
103 |
+
β βββ ...
|
104 |
+
βββ π 256x256.txt π # Contains paths to all 256Γ256 videos
|
105 |
+
β βββ /path/to/video_3.mp4
|
106 |
+
β βββ /path/to/video_4.mp4
|
107 |
+
β βββ ...
|
108 |
+
βββ π 352x288.txt π # Contains paths to all 352x288 videos
|
109 |
+
β βββ /path/to/video_5.mp4
|
110 |
+
β βββ /path/to/video_6.mp4
|
111 |
+
β βββ ...
|
112 |
+
```
|
113 |
+
π Each text file lists video paths corresponding to a specific resolution. Set `args.train_datalist` to the folder containing these files.
|
114 |
+
|
115 |
+
|
116 |
+
---
|
117 |
+
## π **License**
|
118 |
+
|
119 |
+
This project is released under the **MIT License**. See the `LICENSE` file for details.
|
120 |
+
|
121 |
+
|
122 |
+
## π₯ **Why Choose LeanVAE?**
|
123 |
+
LeanVAE is **fast, lightweight and powerful**, enabling high-quality video compression and generation with minimal computational cost.
|
124 |
+
|
125 |
+
If you find this work useful, consider **starring β the repository** and citing our paper!
|
126 |
+
|
127 |
+
---
|
128 |
+
|
129 |
+
## π **Cite Us**
|
130 |
+
```bibtex
|
131 |
+
@misc{cheng2025leanvaeultraefficientreconstructionvae,
|
132 |
+
title={LeanVAE: An Ultra-Efficient Reconstruction VAE for Video Diffusion Models},
|
133 |
+
author={Yu Cheng and Fajie Yuan},
|
134 |
+
year={2025},
|
135 |
+
eprint={2503.14325},
|
136 |
+
archivePrefix={arXiv},
|
137 |
+
primaryClass={cs.CV},
|
138 |
+
url={https://arxiv.org/abs/2503.14325},
|
139 |
+
}
|
140 |
+
```
|
141 |
+
---
|
142 |
+
|
143 |
+
## π **Acknowledgement**
|
144 |
+
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!
|