Home

Awesome

Latte: Latent Diffusion Transformer for Video Generation<br><sub>Official PyTorch Implementation</sub>

<!-- ### [Paper](https://arxiv.org/abs/2401.03048v1) | [Project Page](https://maxin-cn.github.io/latte_project/) --> <!-- [![arXiv](https://img.shields.io/badge/arXiv-2401.03048-b31b1b.svg)](https://arxiv.org/abs/2401.03048) -->

arXiv Project Page HF Demo Hugging Face Spaces slack badge

Static Badge Static Badge

This repo contains PyTorch model definitions, pre-trained weights, training/sampling code and evaluation code for our paper Latte: Latent Diffusion Transformer for Video Generation.

Latte: Latent Diffusion Transformer for Video Generation<br> Xin Ma, Yaohui Wang*, Xinyuan Chen, Gengyun Jia, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, Yu Qiao (*Corresponding Author & Project Lead)

<!-- > <br>Monash University, Shanghai Artificial Intelligence Laboratory,<br> NJUPT, S-Lab, Nanyang Technological University We propose a novel Latent Diffusion Transformer, namely Latte, for video generation. Latte first extracts spatio-temporal tokens from input videos and then adopts a series of Transformer blocks to model video distribution in the latent space. In order to model a substantial number of tokens extracted from videos, four efficient variants are introduced from the perspective of decomposing the spatial and temporal dimensions of input videos. To improve the quality of generated videos, we determine the best practices of Latte through rigorous experimental analysis, including video clip patch embedding, model variants, timestep-class information injection, temporal positional embedding, and learning strategies. Our comprehensive evaluation demonstrates that Latte achieves state-of-the-art performance across four standard video generation datasets, i.e., FaceForensics, SkyTimelapse, UCF101, and Taichi-HD. In addition, we extend Latte to text-to-video generation (T2V) task, where Latte achieves comparable results compared to recent T2V models. We strongly believe that Latte provides valuable insights for future research on incorporating Transformers into diffusion models for video generation. ![The architecture of Latte](visuals/architecture.svg){width=20} --> <!-- <div align="center"> <img src="visuals/architecture.svg" width="650"> </div> This repository contains: * 🪐 A simple PyTorch [implementation](models/latte.py) 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](https://huggingface.co/maxin-cn/Latte/tree/main). * 🛸 A Latte [training script](train.py) using PyTorch DDP. -->

<video controls loop src="https://github.com/Vchitect/Latte/assets/7929326/a650cd84-2378-4303-822b-56a441e1733b" type="video/mp4"></video>

News

# Please update the version of diffusers at leaset to 0.30.0
from diffusers import LattePipeline
from diffusers.models import AutoencoderKLTemporalDecoder
from torchvision.utils import save_image
import torch
import imageio

torch.manual_seed(0)

device = "cuda" if torch.cuda.is_available() else "cpu"
video_length = 16 # 1 (text-to-image) or 16 (text-to-video)
pipe = LattePipeline.from_pretrained("maxin-cn/Latte-1", torch_dtype=torch.float16).to(device)

# Using temporal decoder of VAE
vae = AutoencoderKLTemporalDecoder.from_pretrained("maxin-cn/Latte-1", subfolder="vae_temporal_decoder", torch_dtype=torch.float16).to(device)
pipe.vae = vae

prompt = "a cat wearing sunglasses and working as a lifeguard at pool."
videos = pipe(prompt, video_length=video_length, output_type='pt').frames.cpu()
<!-- <div align="center"> <img src="visuals/latteT2V.gif" width=88%> </div> -->

Setup

First, download and set up the repo:

git clone https://github.com/Vchitect/Latte
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

Sampling

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, just run bash sample/t2v.sh. All related checkpoints will download automatically.

If you would like to measure the quantitative metrics of your generated results, please refer to here.

Training

We provide a training script for Latte in train.py. The structure of the datasets can be found here. 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, these scripts can be also used to train class-conditional and unconditional Latte models. For example, if you want to train the 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

If you are familiar with PyTorch Lightning, you can also use the training script train_pl.py and train_with_img_pl.py provided by @zhang.haojie,

python train_pl.py --config ./configs/ffs/ffs_train.yaml

or

python train_with_img_pl.py --config ./configs/ffs/ffs_img_train.yaml

This script automatically detects available GPUs and uses distributed training.

Contact Us

Yaohui Wang: wangyaohui@pjlab.org.cn Xin Ma: xin.ma1@monash.edu

Citation

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}
}

Acknowledgments

Latte has been greatly inspired by the following amazing works and teams: DiT and PixArt-α, we thank all the contributors for open-sourcing.

License

The code and model weights are licensed under LICENSE.