时间无关的VQGAN和时间敏感的Transformer实现长视频生成 (ECCV 2022)

简介我们提出了TATS,一个长视频生成框架。它在仅包含数十帧的视频上训练,但能够通过滑动窗口生成数千帧的视频。

[最新消息] 我们分析了FVD与人类判断不一致的案例。查看我们的项目网页和论文了解更多信息!

环境配置

  conda create -n tats python=3.8
  conda activate tats
  conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
  pip install pytorch-lightning==1.5.4
  pip install einops ftfy h5py imageio imageio-ffmpeg regex scikit-video tqdm

数据集和预训练模型

UCF-101: 官方数据, VQGAN, TATS-base, TATS-base-uncond
Sky-Timelapse: 官方数据, VQGAN, TATS-base
Taichi-HD: 官方数据, VQGAN, TATS-base
MUGEN: 官方数据, VQGAN, TATS-base
AudioSet-Drums: 官方数据, Video-VQGAN, STFT-VQGAN, TATS-base

生成

短视频: 要生成与训练数据长度相同的视频,请使用scripts/下的代码,并设置以下参数:

gpt_ckpt: 训练好的transformer检查点路径。
vqgan_ckpt: 训练好的VQGAN检查点路径。
save: 保存生成结果的路径。
save_videos: 表示将保存视频。
class_cond: 表示使用类别标签作为条件信息。

要计算FVD,需要以下参数:

compute_fvd: 表示将计算FVD。
data_path: 数据集文件夹路径。
dataset: 数据集名称。
image_folder: 当数据集包含帧而不是视频时使用,例如Sky Time-lapse。
sample_every_n_frames: 在真实视频数据中跳过的帧数,例如使用Taichi-HD数据集时请设置为4。
resolution: 计算FVD时真实视频的分辨率,例如UCF、Sky和Taichi为128,MUGEN为256。

python sample_vqgan_transformer_short_videos.py \
    --gpt_ckpt {GPT-CKPT} --vqgan_ckpt {VQGAN-CKPT} --class_cond \
    --save {SAVEPATH} --data_path {DATAPATH} --batch_size 16 --resolution 128 \
    --top_k 2048 --top_p 0.8 --dataset {DATANAME} --compute_fvd --save_videos

长视频: 要使用滑动窗口生成长于训练长度的视频,请使用以下脚本。

sample_length: 要生成的潜在帧数。
temporal_sample_pos: 滑动窗口方法生成的帧位置。

python sample_vqgan_transformer_long_videos.py \
    --gpt_ckpt {GPT-CKPT} --vqgan_ckpt {VQGAN-CKPT} \
    --dataset ucf101 --class_cond --sample_length 16 --temporal_sample_pos 1 --batch_size 5 --n_sample 5 --save_videos

文本生成视频: 要根据文本条件生成MUGEN视频,请查看这个Colab笔记本示例!
音频生成视频: 要根据音频条件生成鼓声视频,请使用以下脚本。

stft_vqgan_ckpt: 用于STFT特征的训练好的VQGAN检查点路径。

python sample_vqgan_transformer_audio_cond.py \
    --gpt_ckpt {GPT-CKPT} --vqgan_ckpt {VQGAN-CKPT} --stft_vqgan_ckpt {STFT-CKPT} \
    --dataset drum --n_sample 10

分层采样: 要使用分层模型(先用AR transformer,再用插值transformer)生成长于训练长度的视频,请使用以下检查点和脚本。

Sky检查点: VQGAN, AR-transformer, Interpolation-transformer
Taichi检查点: VQGAN, AR-transformer, Interpolation-transformer
top_k_init, top_p_init: 这些参数允许在采样初始帧时使用不同的top_k和top_p值。例如,在采样天空视频时可以使用--top_k_init 16384 --top_p_init 0.99来增加多样性。

python sample_vqgan_transformer_hierarchical.py \
    --ckpt1 {AR-CKPT} --CKPT2 {Interpolation-CKPT} --vqgan {VQGAN-CKPT} \
    --dataset sky --top_k_init 2048 --top_p_init 0.8 --top_k 2048 --top_p 0.8 --temporal_sample_pos 1

训练

以下是训练VQGAN和transformer的示例用法。根据不同设置可能需要更改的参数说明:

data_path: 数据集文件夹路径。
default_root_dir: 保存检查点和tensorboard日志的路径。
vqvae: 训练好的VQGAN检查点路径。
resolution: 训练视频片段的分辨率。
sequence_length: 训练视频片段的帧数。
discriminator_iter_start: 开始使用GAN损失的步骤ID。
image_folder: 当数据集包含帧而不是视频时使用，例如天空延时摄影。
unconditional: 当没有条件信息可用时使用此标志，例如天空延时摄影。
sample_every_n_frames: 在真实视频数据中跳过的帧数，例如在Taichi-HD数据集上训练时，请将其设置为4。
downsample: 时间、高度和宽度维度的采样率。
no_random_restart: 是否重新初始化码书标记。

VQGAN

python train_vqgan.py --embedding_dim 256 --n_codes 16384 --n_hiddens 32 --downsample 4 8 8 --no_random_restart \
                      --gpus 8 --sync_batchnorm --batch_size 2 --num_workers 32 --accumulate_grad_batches 6 \
                      --progress_bar_refresh_rate 500 --max_steps 2000000 --gradient_clip_val 1.0 --lr 3e-5 \
                      --data_path {DATAPATH} --default_root_dir {CKPTPATH} \
                      --resolution 128 --sequence_length 16 --discriminator_iter_start 10000 --norm_type batch \
                      --perceptual_weight 4 --image_gan_weight 1 --video_gan_weight 1  --gan_feat_weight 4

Transformer

TATS-base Transformer

python train_transformer.py --num_workers 32 --val_check_interval 0.5 --progress_bar_refresh_rate 500 \
                        --gpus 8 --sync_batchnorm --batch_size 3 --unconditional \
                        --vqvae {VQGAN-CKPT} --data_path {DATAPATH} --default_root_dir {CKPTPATH} \
                        --vocab_size 16384 --block_size 1024 --n_layer 24 --n_head 16 --n_embd 1024  \
                        --resolution 128 --sequence_length 16 --max_steps 2000000

要训练条件Transformer，请移除--unconditional标志并使用以下标志

cond_stage_key: 要使用的条件信息类型。可以是label、text或stft。
stft_vqvae: 用于STFT特征的训练好的VQGAN检查点路径。
text_cond: 使用此标志表示BPE编码的文本。

TATS-hierarchical Transformer

python train_transformer.py --num_workers 32 --val_check_interval 0.5 --progress_bar_refresh_rate 500 \
                        --gpus 8 --sync_batchnorm --batch_size 3 --unconditional \
                        --vqvae {VQGAN-CKPT} --data_path {DATAPATH} --default_root_dir {CKPTPATH} \
                        --vocab_size 16384 --block_size 1280 --n_layer 24 --n_head 16 --n_embd 1024  \
                        --resolution 128 --sequence_length 20 --spatial_length 128 --n_unmasked 256 --max_steps 2000000

python train_transformer.py --num_workers 32 --val_check_interval 0.5 --progress_bar_refresh_rate 500 \
                        --gpus 8 --sync_batchnorm --batch_size 4 --unconditional \
                        --vqvae {VQGAN-CKPT} --data_path {DATAPATH} --default_root_dir {CKPTPATH} \
                        --vocab_size 16384 --block_size 1024 --n_layer 24 --n_head 16 --n_embd 1024  \
                        --resolution 128 --sequence_length 64 --sample_every_n_latent_frames 4 --spatial_length 128 --max_steps 2000000

致谢

我们的代码部分基于VQGAN和VideoGPT。

引用

@article{ge2022long,
  title={Long video generation with time-agnostic vqgan and time-sensitive transformer},
  author={Ge, Songwei and Hayes, Thomas and Yang, Harry and Yin, Xi and Pang, Guan and Jacobs, David and Huang, Jia-Bin and Parikh, Devi},
  journal={arXiv preprint arXiv:2204.03638},
  year={2022}
}

许可证

TATS采用MIT许可证，详见LICENSE文件。