nielsr HF staff commited on
Commit
9fa6488
Β·
verified Β·
1 Parent(s): dd19028

Add pipeline tag to model card

Browse files

This PR adds the `pipeline_tag: image-to-image` to the model card metadata. This makes the model discoverable through the model hub interface.

Files changed (1) hide show
  1. README.md +144 -3
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!