Sora - Open-Sora技术报告1.0

vincent&lin

于 2024-08-12 12:01:40 发布

阅读量31

点赞数

文章标签： ai

原文链接：https://github.com/hpcaitech/Open-Sora/blob/v1.2.0/docs/report_01.md

版权

源：Open-Sora/docs/report_01.md at v1.2.0 · hpcaitech/Open-Sora · GitHub

OpenAI's Sora能够生成1分钟的高质量视频，但是并没有对外批露任何实现细节。为了让AI更"open"，我们致力于构建一个开源版本的Sora。这个报告描述了我们训练一个基于Transformer的视频扩散模型的初次尝试。

架构选型

为了降低算力开销，我们希望利用已有的VAE模型，Sora使用时空VAE模型降低时间维度。但是我们并没有找到高质量的开源时空VAE模型。 MAGVIT的4x4x4 VAE没有开源，VideoGPT的2x4x4 VAE质量较差。最终我们决定在第一个版本中使用Stability-AI的2D VAE模型。

视频模型训练包含了大量的tokens。考虑帧率24fps的1分钟视频，有1440帧。使用VAE下采样4x，patch size下采样2x，我们就拥有1440x1024≈1.5M tokens（假设原始尺寸是 2048×2048像素，下采样4倍后变为 512×512 像素，再用16×16 的patch下采样原尺寸中的 32×32 的patch，可以得到1024 tokens/frame）。

在1.5M tokens上做Full Attention需要超大的算力开销。所以我们使用Latte的spatial-temporal attention去降低算力开销。

如图所示，在STDiT中，我们在每个spatial attention后面插入1个temporal attention，类似Latte论文中的变体3。但是我们的参数量和Latte有差异，Latte的论文中声称他们的变体要优于变体3，但是我们在16x256x256的视频上的实验得到的性能排名：DiT (full) > STDiT (Sequential) > STDiT (Parallel) ≈ Latte。所以我们最终从中选择了STDiT(Sequential)。

为了专注视频生成，我们希望基于一个强大的图像生成模型来训练模型。 PixArt-α 是一个T5-conditioned DiT结构的，训练效率高的图像生成模型。我们使用PixArt-α来初始化我们的模型，并用0初始化插入的temporal attention的投影层。这种初始化在开始阶段保留了模型的图像生成能力，而Latte架构不能。这个插入的attention将参数量从580M提升到了724M。

借鉴PixArt-α和Stable Video Diffusion，我们采用了激进的训练策略。在366K数据集上进行16x256x256的预训练, 再在20K数据集上进行16x256x256, 16x512x512, 和64x512x512预训练。通过可伸缩的position embedding，这个策略大幅降低了算力开销。

我们还在DiT中尝试了3D patch embedder，在temporal维度使用2x下采样，生成的视频质量较差。所以我们将temporal下采样放到了下一个版本。当前我们在16帧训练中每3帧采样1次，在64帧训练中，每2帧采样1次。

Data is the key to high quality

我们发现数据的质量和数量对生成视频的质量有很大影响，甚至超过了模型结果和训练策略。当前我们使用了 HD-VG-130M的366K的视频clips，这些视频的质量参差不齐，说明文字也不准确。所以我们又从Pexels中收集了20K相对高质量的视频。我们使用图像说明文字生成模型LLaVA对视频进行标注，每3帧1个prompt。使用设计好的prompt，LLaVA可以生成高质量的说明文字。