【2024】Open-Sora:揭示完整的模型参数、训练细节以及类似 Sora 的视频生成模型的一切
不久前,OpenAI Sora以其惊人的视频生成效果迅速走红,在众多文本转视频模型中脱颖而出,成为全球关注的焦点。继两周前推出用于训练和推理过程的Sora Replication,成本降低46%之后,Colossal-AI团队又推出了新的开源解决方案“Open-Sora 1.0”,涵盖了整个训练过程,包括数据处理、所有训练细节和模型检查点,与世界各地的AI爱好者携手推进视频创作的新时代。
我们先来看看Colossal-AI团队发布的“Open-Sora 1.0”模型生成的繁华都市夜景视频。
Open-Sora 1.0 生成的城市景观
这只是Open-Sora技术的冰山一角。Colossal-AI团队在我们的GitHub上完全开源了模型架构、 训练后的模型检查点、所有训练和数据准备过程的详细信息、视频演示和教程, 供所有对文本转视频模型感兴趣的人免费学习和使用。我们将持续更新Open-Sora相关解决方案和最新动态。我们将持续更新 Open-Sora 解决方案和此 GitHub 存储库中的最新消息,欢迎大家关注更新!
- Open-Sora开源地址:****https://github.com/hpcaitech/Open-Sora
Open Sora全面讲解
接下来,我们将对Sora复现方案的多个关键维度进行深入讲解,包括模型架构设计、训练复现方案、数据预处理、模型输出演示以及高效的训练优化策略。
模型架构设计
我们的模型采用当前流行的扩散变压器(DiT)[1]架构。我们使用 PixArt-α [2],这是一种高质量的开源文本到图像模型,也使用 DiT 架构作为基础,并通过添加时间注意力层将其扩展为生成视频。具体来说,整个架构由预训练的 VAE、文本编码器和利用时空注意力机制的 STDiT(时空扩散变换器)模型组成。STDiT各层的结构如下所示。它使用串行方法将 1D 时间注意力模块叠加在 2D 空间注意力模块上,以对时间关系进行建模。在时间注意力模块之后,使用交叉注意力模块来对齐文本的语义。与全注意力机制相比,这样的结构大大降低了训练和推理成本。与同样采用时空注意力机制的 Latte [3] 模型相比,STDiT 可以更好地利用预训练图像 DiT 的权重来继续对视频数据进行训练。
STDiT模型结构示意图
整个模型的训练和推理过程如下。在训练阶段,我们首先使用预训练的 VAE(变分自动编码器)编码器来压缩视频数据,然后在压缩后在潜在空间中使用文本嵌入来训练我们提出的 STDiT 模型。在推理阶段,我们从VAE的潜在空间中随机采样一个高斯噪声,并将其与提示嵌入一起输入到STDiT中,以获得去噪后的特征,最后输入到VAE解码器中以获得视频。
模型训练过程
训练复制方案
Open-Sora 的再现方案参考了稳定视频扩散 (SVD) [3] 工作,由三个阶段组成,即:
每个阶段都根据前一阶段的权重继续训练。与从头开始的单阶段训练相比,多阶段训练通过逐步扩大数据,更高效地实现高质量视频生成的目标。
- 第一阶段:大规模图像预训练
第一阶段借助成熟的Text-to-Image模型,通过大规模图像预训练,有效降低视频预训练的成本。
借助互联网上丰富的大规模图像数据和先进的Text-to-Image技术,我们可以训练出高质量的Text-to-Image模型,作为下一阶段视频预处理的初始化权重。训练。同时,由于没有高质量的时空VAE,我们利用稳定扩散[5]模型来预训练空间VAE。该策略不仅保证了初始模型的优越性能,而且显着降低了视频预训练的总体成本。
- 第二阶段:大规模视频预训练
第二阶段进行大规模视频预训练,增加模型的泛化能力,有效掌握视频的时间序列关联性。
该阶段需要使用大量的视频数据训练,保证视频主题的多样性,从而增加模型的泛化能力。第二阶段的模型在第一阶段的 Test-to-Image 模型中添加了时间序列注意模块,用于学习视频中的时间关系。其余模块与第一阶段保持相同,并加载第一阶段权重作为初始化,同时将时间注意模块的输出初始化为零,以实现更高效、更快的收敛。我们使用PixArt-α[2]的开源权重作为第二阶段STDiT模型的初始化,并采用T5[6]模型作为文本编码器。同时,我们使用256x256的小分辨率进行预训练,进一步提高了收敛速度并降低了训练成本。
- 第三阶段:高质量视频数据微调
第三阶段对高质量视频数据进行微调,显着提高生成视频的质量。
第三阶段使用的视频数据大小比第二阶段小一个数量级,但视频的持续时间、分辨率和质量更高。通过这种方式进行微调,我们实现了视频生成从短到长、从低到高分辨率、从低到高保真度的有效缩放。
在我们的重现过程中,我们使用了 64 个 H800 GPU 进行训练。第二阶段的训练量总计为 2,808 GPU 小时,约合 7,000 美元,第三阶段的训练量为 1,920 GPU 小时,约合 4,500 美元,我们成功地将Open-Sora 复现过程保持在约 10,000 美元。
数据预处理
为了进一步降低Sora复制的门槛和复杂性,Colossal-AI团队还在代码库中提供了便捷的视频数据预处理脚本,包括公共视频数据集下载、根据镜头连续性将长视频分割为短视频片段、使用开放的源码大语言模型LLaVA[7]生成细粒度的提示词,这样就可以轻松开始Sora复制预训练。我们提供的批量视频字幕生成代码可以在 3 秒内使用两个 GPU 对视频进行注释,质量接近 GPT-4V。最终的视频文本对可以直接用于训练。通过我们提供的代码,用户可以轻松快速地生成在自己的数据集上进行训练所需的视频文本对,从而显着降低技术障碍并为启动 Sora 复制项目做好准备。
基于数据预处理脚本自动生成视频/文本对
视频演示
我们来看看Open-Sora的实际视频生成结果。例如,让 Open-Sora 生成大海拍打悬崖海岸岩石的航拍镜头。
然后让我们的 Open-Sora 鸟瞰壮丽的山瀑布从悬崖上奔流而下,最终汇入湖中的壮丽景色。
除了天空之外,我们还可以生成有关海洋的视频。通过简单的输入提示,我们让 Open-Sore 生成一段关于海龟在珊瑚礁中悠闲游泳的视频。
Open-Sora 还能够通过延时摄影向我们展示银河系的繁星璀璨。
如果您对视频生成有更多有趣的想法,欢迎来到我们的Open-Sora开源社区获取模型权重免费体验。
链接:https: //github.com/hpcaitech/Open-Sora
值得注意的是,我们当前的版本仅使用了 400K 的训练数据,模型的生成质量和文本跟踪能力还有待提高。例如,在上面的乌龟视频中,生成的乌龟有一只额外的脚。Open-Sora 1.0 也不太擅长生成肖像和复杂图像。我们在 GitHub 上有一份计划列表,将继续解决现有的 bug 并提高生成的质量。
高效培训
除了大幅降低Sora复现的技术门槛,在时长、分辨率、内容等多个维度提升视频生成质量外,我们还提供Colossal-AI加速系统,实现高效训练。通过内核优化和混合并行等高效训练策略,我们在处理 64 帧、512x512 分辨率视频的训练中实现了1.55 倍的加速。同时,得益于Colossal-AI的异构内存管理系统,我们可以毫无阻碍地在单台服务器(8H800)上执行1分钟1080p高清视频训练任务。*
此外,我们提出的 STDiT 模型架构在训练过程中表现出了卓越的效率。与具有全注意力机制的 DiT 相比, STDiT随着帧数的增加实现了高达 5 倍的加速,这在处理长视频序列等现实任务中尤其关键。
Open-Sora 视频一览
欢迎持续关注Open-Sora开源项目:https://github.com/hpcaitech/Open-Sora
后续他们将持续维护和优化Open-Sora项目,利用更多的视频训练数据生成更高质量、更长的视频内容,并支持多分辨率特性,有效推动AI技术在电影、游戏、广告等领域的应用。
参考文章:【2024】Open-Sora:揭示完整的模型参数、训练细节以及类似 Sora 的视频生成模型的一切
参考链接:
[1] https://arxiv.org/abs/2212.09748 Scalable Diffusion Models with Transformers
[2] https://arxiv.org/abs/2310.00426 PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis
[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] https://huggingface.co/stabilityai/sd-vae-ft-mse-original
[6] https://github.com/google-research/text-to-text-transfer-transformer