This repo contains PyTorch model definitions, pre-trained weights, and training/sampling code for our paper exploring latent diffusion models with transformers (Latte). You can find more visualizations on our project page.
Latte: Latent Diffusion Transformer for Video Generation
Xin Ma, Yaohui Wang*, Xinyuan Chen, Gengyun Jia, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, Yu Qiao (*Corresponding Author & Project Lead)
This repository contains:
-
๐ช A simple PyTorch implementation of Latte
-
โก๏ธ Pre-trained Latte models trained on FaceForensics, SkyTimelapse, Taichi-HD and UCF101 (256x256). In addition, we provide a T2V checkpoint (512x512). All checkpoints can be found here.
-
๐ธ A Latte training script using PyTorch DDP.
- (๐ฅ New) Mar. 20, 2024. ๐ฅ An updated LatteT2V model is coming soon, stay tuned!
-
(๐ฅ New) Feb. 24, 2024. ๐ฅ We are very grateful that researchers and developers like our work. We will continue to update our LatteT2V model, hoping that our efforts can help the community develop. Our Latte discord channel
is created for discussions. Coders are welcome to contribute.
-
(๐ฅ New) Jan. 9, 2024. ๐ฅ A updated LatteT2V model initialized with the PixArt-ฮฑ is released, checkpoint can be found here.
-
(๐ฅ New) Oct. 31, 2023. ๐ฅ The training and inference code is released. All checkpoints (including FaceForensics, SkyTimelapse, UCF101, and Taichi-HD) can be found here. In addition, LatteT2V inference code is provided.
First, download and set up the repo:
git clone https://github.com/maxin-cn/Latte.git
cd Latte
We provide an environment.yml
file that can be used to create a Conda environment. If you only want
to run pre-trained models locally on CPU, you can remove the cudatoolkit
and pytorch-cuda
requirements from the file.
conda env create -f environment.yml
conda activate latte
You can sample from our pre-trained Latte models with sample.py
. Weights for our pre-trained Latte model can be found here. The script has various arguments to adjust sampling steps, change the classifier-free guidance scale, etc. For example, to sample from our model on FaceForensics, you can use:
bash sample/ffs.sh
or if you want to sample hundreds of videos, you can use the following script with Pytorch DDP:
bash sample/ffs_ddp.sh
If you want to try generating videos from text, please download t2v_required_models
and run bash sample/t2v.sh
.
We provide a training script for Latte in train.py
. This script can be used to train class-conditional and unconditional
Latte models. To launch Latte (256x256) training with N
GPUs on the FaceForensics dataset
:
torchrun --nnodes=1 --nproc_per_node=N train.py --config ./configs/ffs/ffs_train.yaml
or If you have a cluster that uses slurm, you can also train Latte's model using the following scripts:
sbatch slurm_scripts/ffs.slurm
We also provide the video-image joint training scripts train_with_img.py
. Similar to train.py
scripts, this scripts can be also used to train class-conditional and unconditional
Latte models. For example, if you wan to train Latte model on the FaceForensics dataset, you can use:
torchrun --nnodes=1 --nproc_per_node=N train_with_img.py --config ./configs/ffs/ffs_img_train.yaml
Yaohui Wang: [email protected] Xin Ma: [email protected]
If you find this work useful for your research, please consider citing it.
@article{ma2024latte,
title={Latte: Latent Diffusion Transformer for Video Generation},
author={Ma, Xin and Wang, Yaohui and Jia, Gengyun and Chen, Xinyuan and Liu, Ziwei and Li, Yuan-Fang and Chen, Cunjian and Qiao, Yu},
journal={arXiv preprint arXiv:2401.03048},
year={2024}
}
Latte has been greatly inspired by the following amazing works and teams: DiT and PixArt-ฮฑ, we thank all the contributors for open-sourcing.
The code and model weights are licensed under LICENSE.