OpenAI Sora尚未使用，国产Sora刚刚正式宣布全面开源！「Open-Sora 1.0」所有的训练细节模型权重全面开源，使用1万美元的64块GPU进行复现，训练成本直降46%。
最近，OpenAI Sora以其令人惊叹的视频生成效果迅速走红，在众多文生视频模型中脱颖而出，成为全球瞩目的焦点。
在2周前，Colossal-AI团队推出了一款名为Sora的训练推理复现流程，成功将成本降低了46%。现在，他们全面开源了全球首个类Sora架构视频生成模型「Open-Sora 1.0」。这个模型涵盖了整个训练流程，包括数据处理、所有训练细节和模型权重。Colossal-AI团队希望与全球的AI热爱者一起推进视频创作的新纪元。
Open-Sora的开源代码可以在以下地址找到：https://github.com/hpcaitech/Open-Sora
让我们迅速来看一段由Colossal-AI团队发布的「Open-Sora 1.0」模型生成的都市繁华掠影视频。
Open-Sora 1.0创造的都市繁华景象
这只是Sora复现技术冰山的一小部分，关于以上视频的模型结构、训练好的模型参数、复现的所有训练细节、数据预处理过程、演示和详细的入门教程，Colossal-AI团队已经完全免费开源在GitHub上。
同时，我们立即联系了该团队，了解到他们将持续更新Open-Sora的相关解决方案和最新动态，对此感兴趣的朋友可以持续关注Open-Sora的开源社区。
详细解析Sora复现方案
下面，我们将详细讨论Sora复现方案的几个重要方面，包括模型架构设计、训练复现方案、数据预处理、模型生成效果展示以及高效训练优化策略。
模型架构设计
该模型采用了当前流行的Diffusion Transformer（DiT）[1]架构。
作者团队在使用DiT架构的高质量开源文生图模型PixArt-α [2]的基础上，通过引入时间注意力层，将其扩展到了视频数据上。
具体来说，整个架构包括一个经过预训练的变分自编码器（VAE），一个文本编码器，以及一个利用空间-时间注意力机制的STDiT（Spatial Temporal Diffusion Transformer）模型。
其中，STDiT 每层的结构如下图所示。它采用串行的方式在二维的空间注意力模块上叠加一维的时间注意力模块，用于建模时序关系。
在时间注意力模块之后，交叉注意力模块用于对齐文本的语义。与全注意力机制相比，这样的结构显著减少了训练和推理的成本。
与同样采用空间-时间注意力机制的Latte [3]模型相比，STDiT能更充分地利用已经预训练好的图像DiT的权重，在视频数据上进行进一步训练。
STDiT结构示意图
模型的训练过程
复现训练方案
我们从该团队了解到，Open-Sora的复现方案参考了稳定视频扩散（SVD）[3]的工作，总共包括三个阶段，分别是：
1. 图像的大规模预训练；
2. 视频预训练的大规模应用；
3. 对高质量视频数据进行微调。
每个阶段都会基于前一个阶段的权重继续训练。与从头开始单阶段训练相比，多阶段训练通过逐步扩展数据，更有效地实现高质量视频生成的目标。
培训计划的三个阶段
第一阶段：图像的大规模预训练
首先，我们可以利用成熟的视觉生成模型，通过大规模图像预训练来降低视频预训练的成本。
根据作者团队的透露，利用互联网上丰富的大规模图像数据和先进的文生图技术，我们可以训练一个优秀的文生图模型，该模型将作为下一阶段视频预训练的初始权重。
然而，由于目前没有高质量的时空VAE可用，他们选择了使用经过Stable Diffusion[5]模型预训练的图像VAE。这种策略不仅确保了初始模型的出色性能，还大幅降低了视频预训练的总体成本。
第二阶段：视频预训练的大规模实施
在第二阶段，我们将进行大规模视频预训练，以提高模型的泛化能力，并有效地捕捉视频中的时间序列关联。
据我们了解，这个阶段需要使用大量视频数据进行训练，以确保视频题材的多样性，从而提高模型的泛化能力。第二阶段的模型在第一阶段生成图像模型的基础上，引入了时序注意力模块，用于学习视频中的时序关系。
其他模块与第一阶段保持一致，并使用第一阶段的权重进行初始化。同时，在初始化时，注意力模块的输出设置为零，以实现更高效和更快速的收敛。
Colossal-AI团队在第二阶段STDiT模型中使用了PixArt-alpha[2]的开源权重作为初始化，并选择了T5 [6]模型作为文本编码器。此外，他们还采用了256×256的低分辨率进行预训练，以加快收敛速度并降低训练成本。
第三阶段：优化视频数据质量
第三阶段对高质量视频数据进行微调，显著提升视频生成的质量。
作者团队提到，在第三阶段使用的视频数据规模比第二阶段少了一个数量级，但是视频的时长、分辨率和质量都更高。通过这种方式进行微调，他们成功地实现了视频生成的高效扩展，从短视频到长视频、从低分辨率到高分辨率、从低保真度到高保真度。
据作者团队透露，在Open-Sora的复现过程中，他们采用了64块H800进行训练。
第二阶段的训练总时长为2808个GPU小时，相当于约7000美元。第三阶段的训练总时长为1920个GPU小时，大约为4500美元。经过初步估算，整个训练方案成功将Open-Sora复现流程的成本控制在了1万美元左右。
数据预处理
为了进一步降低Sora复现的门槛和复杂度，Colossal-AI团队在代码仓库中还提供了方便的视频数据预处理脚本，让用户可以轻松启动Sora复现预训练。这些脚本包括公开视频数据集的下载，将长视频根据镜头连续性分割为短视频片段，并使用开源大语言模型LLaVA [7]生成精细的提示词。
作者团队提到他们提供的批量视频标题生成代码可以在两卡的情况下，以3秒的速度为一个视频生成标注，并且生成的质量接近于GPT-4V。最终得到的视频和文本对可以直接用于训练。
通过利用他们在GitHub上提供的开源代码，我们可以方便地在自己的数据集上快速生成训练所需的视频/文本对，大大降低了开始Sora复现项目所需的技术门槛和前期准备工作。
通过使用数据预处理脚本，我们可以自动生成视频和文本的配对。
展示模型生成的效果
接下来，我们将展示Open-Sora的实际视频生成效果。例如，让Open-Sora生成一段航拍画面，拍摄悬崖海岸边海水拍打着岩石的场景。
让Open-Sora去捕捉山川瀑布从悬崖上壮丽地倾泻而下，最终汇入湖泊的壮观鸟瞰画面。
除了上天还能入海，只需简单使用prompt函数，让Open-Sora生成了一段水中世界的场景，场景中一只海龟在珊瑚礁间自由游动。
Open-Sora还能通过延时摄影的方式，向我们展示了闪烁的星空。
如果你还有更多视频生成的有趣想法，可以访问Open-Sora开源社区获取模型权重进行免费的尝试。
请点击以下链接查看：https://github.com/hpcaitech/Open-Sora
需要注意的是，开发团队在Github上指出，目前版本仅使用了400K的训练数据，模型的生成质量和对文本的遵循能力还有待提高。例如，在上述乌龟视频中，生成的乌龟多了一只脚。Open-Sora 1.0也不擅长生成人像和复杂画面。
作者团队在Github上列举了一系列计划，旨在持续改进现有问题，提高生成结果的质量。
提升训练效率
除了显著降低Sora复现的技术门槛，提升视频生成在时长、分辨率、内容等多个方面的质量，作者团队还提供了Colossal-AI加速系统来加快Sora复现的训练过程。
通过算法优化和混合并行等高效训练策略，在处理64帧、512×512分辨率视频的训练中，实现了1.55倍的加速效果。
同时，由于Colossal-AI的异构内存管理系统，单台服务器（8 x H800）可以轻松完成1分钟的1080p高清视频训练任务。
此外，在作者团队的研究中，我们还发现STDiT模型在训练过程中表现出了出色的高效性。
与使用全注意力机制的DiT相比，STDiT在增加帧数的情况下实现了高达5倍的加速效果，这对于处理长视频序列等实际任务非常重要。
查看Open-Sora模型视频生成的效果
欢迎继续关注Open-Sora开源项目：https://github.com/hpcaitech/Open-Sora
根据作者团队的说法，他们计划继续改进和优化Open-Sora项目。他们打算增加更多的视频训练数据，以生成更高质量、更长时长的视频内容，并支持多分辨率特性。这将有助于推动AI技术在电影、游戏、广告等领域的应用。
参考文献：
[1] 在这篇论文中，作者介绍了一种使用Transformer的可扩展扩散模型。
[2] https://arxiv.org/abs/2310.00426 PixArt-α: Efficient Training of Diffusion Transformer for Realistic Text-to-Image Generation
[3] https://arxiv.org/abs/2311.15127 Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets
[4] https://arxiv.org/abs/2401.03048 Latte: Latent Diffusion Transformer for Video Generation
[5] You can find the original code at the following link: https://huggingface.co/stabilityai/sd-vae-ft-mse-original
[6] You can find the code for the text-to-text transfer transformer model on GitHub at the following link: https://github.com/google-research/text-to-text-transfer-transformer
[7] 你可以在这里找到LLaVA的GitHub链接：https://github.com/haotian-liu/LLaVA
