Latte

Latte：用于视频生成的潜在扩散变换器
_{官方PyTorch实现}

本仓库包含我们论文《Latte：用于视频生成的潜在扩散变换器》的PyTorch模型定义、预训练权重、训练/采样代码和评估代码。

新闻

(🔥 新) 2024年7月11日 💥 Latte-1现已集成到diffusers中。感谢@yiyixuxu、@sayakpaul、@a-r-r-o-w和@DN6。 您可以使用以下代码轻松运行Latte。我们还支持4/8位量化推理，可将GPU内存从17 GB减少到9 GB。更多信息请参阅此教程。

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 (文本到图像) 或 16 (文本到视频)
pipe = LattePipeline.from_pretrained("maxin-cn/Latte-1", torch_dtype=torch.float16).to(device)

# 使用VAE的时间解码器
vae = AutoencoderKLTemporalDecoder.from_pretrained("maxin-cn/Latte-1", subfolder="vae_temporal_decoder", torch_dtype=torch.float16).to(device)
pipe.vae = vae

prompt = "一只戴着太阳镜在游泳池担任救生员的猫。"
videos = pipe(prompt, video_length=video_length, output_type='pt').frames.cpu()

(🔥 新) 2024年5月23日 💥 Latte-1已发布！预训练模型可在此处下载。我们同时支持T2V和T2I。请分别运行bash sample/t2v.sh和bash sample/t2i.sh。

-->

(🔥 新) 2024年2月24日 💥 我们非常感谢研究人员和开发者喜欢我们的工作。我们将继续更新我们的LatteT2V模型,希望我们的努力能够帮助社区发展。我们的Latte discord频道已创建用于讨论。欢迎程序员贡献。
(🔥 新) 2024年1月9日 💥 发布了一个使用PixArt-α初始化的更新版LatteT2V模型,可以在这里找到检查点。
(🔥 新) 2023年10月31日 💥 训练和推理代码已发布。所有检查点(包括FaceForensics、SkyTimelapse、UCF101和Taichi-HD)可以在这里找到。此外,还提供了LatteT2V推理代码。

设置

首先,下载并设置仓库:

git clone https://github.com/Vchitect/Latte
cd Latte

我们提供了一个environment.yml文件,可用于创建Conda环境。如果您只想在本地CPU上运行预训练模型,可以从文件中删除cudatoolkit和pytorch-cuda要求。

conda env create -f environment.yml
conda activate latte

采样

您可以使用sample.py从我们的预训练Latte模型中采样。我们预训练的Latte模型的权重可以在这里找到。该脚本有各种参数可以调整采样步骤,更改无分类器引导尺度等。例如,要从我们在FaceForensics上的模型中采样,您可以使用:

bash sample/ffs.sh

或者如果您想采样数百个视频,可以使用以下带有Pytorch DDP的脚本:

bash sample/ffs_ddp.sh

如果您想尝试从文本生成视频,只需运行bash sample/t2v.sh。所有相关的检查点将自动下载。

如果您想测量生成结果的定量指标,请参考这里。

训练

我们在train.py中提供了Latte的训练脚本。数据集的结构可以在这里找到。此脚本可用于训练类条件和无条件的Latte模型。要在FaceForensics数据集上使用N个GPU启动Latte(256x256)训练:

torchrun --nnodes=1 --nproc_per_node=N train.py --config ./configs/ffs/ffs_train.yaml

或者如果您有使用slurm的集群,也可以使用以下脚本训练Latte模型:

sbatch slurm_scripts/ffs.slurm

我们还提供了视频-图像联合训练脚本train_with_img.py。与train.py脚本类似,这些脚本也可以用于训练类条件和无条件的Latte模型。例如,如果您想在FaceForensics数据集上训练Latte模型,可以使用:

torchrun --nnodes=1 --nproc_per_node=N train_with_img.py --config ./configs/ffs/ffs_img_train.yaml

如果您熟悉PyTorch Lightning,也可以使用@zhang.haojie提供的训练脚本train_pl.py,

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

此脚本会自动检测可用的GPU并使用分布式训练。

联系我们

王耀辉: wangyaohui@pjlab.org.cn 马鑫: xin.ma1@monash.edu

引用

如果您发现这项工作对您的研究有用,请考虑引用它。

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