全球首個開源的類別Sora架構視訊生成模型，來了！
整個訓練流程，包括資料處理、所有訓練細節和模型權重，全部開放。
這就是剛發布的Open-Sora 1.0。
它帶來的實際效果如下，能生成繁華都市夜景中的車水馬龍。
還能用空拍視角，展現懸崖海岸邊，海水拍打著岩石的畫面。
也或是縮時攝影下的浩瀚星空。
自Sora發布以來，由於其驚人的效果和技術細節的稀缺性，揭示和復現Sora已成為開發社群中最受關注的話題之一。例如，Colossal-AI團隊推出了一項能夠降低46%成本的Sora訓練和推理重現流程。
短短兩週時間後，該團隊再次發布最新進展，復現類Sora方案，並將技術方案及詳細上手教程在GitHub上免費開源。
那麼問題來了，復現Sora如何實現？
Open-Sora 開源位址：https://github.com/hpcaitech/Open-Sora
全面解讀Sora復現方案
Sora復現方案包括四個面向：
- 模型架構設計
- 訓練復現方案
- 資料預處理
- 高效訓練最佳化策略
模型架構設計
模型採用了Sora同源架構Diffusion Transformer (DiT) 。
它以採用DiT架構的高品質開源文生圖模型PixArt-α為基座，在此基礎上引入時間注意力層，將其擴展到視訊資料上。
具體來看，整個架構包括一個預先訓練好的VAE，一個文字編碼器和一個利用空間-時間注意力機制的STDiT (Spatial Temporal Diffusion Transformer)模型。
其中，STDiT 每層的結構如下圖所示。
它採用串列的方式在二維的空間注意力模組上疊加一維的時間注意力模組，用於建模時序關係。在時間注意力模組之後，交叉注意力模組用於對齊文字的語義。
與全注意力機制相比，這樣的結構大大降低了訓練和推理開銷。
與同樣使用空間-時間注意力機制的 Latte模型相比，STDiT 可以更好的利用已經預先訓練好的影像 DiT 的權重，從而在影片資料上繼續訓練。
△STDiT結構示意圖
整個模型的訓練和推理流程如下。
據了解，在訓練階段首先採用預先訓練好的Variational Autoencoder (VAE)的編碼器將視訊資料進行壓縮，然後在壓縮之後的潛在空間中與文字嵌入(text embedding)一起訓練STDiT擴散模型。
在推理階段，從VAE的潛在空間中隨機採樣出一個高斯噪聲，與提示詞嵌入(prompt embedding)一起輸入到STDiT中，得到去噪之後的特徵，最後輸入到VAE的解碼器，解碼得到視頻。
△模型训练流程
训练复现方案
在训练复现部分，Open-Sora参考了Stable Video Diffusion (SVD)。
一共分为3个阶段：
- 大规模图像预训练。
- 大规模视频预训练。
- 高质量视频数据微调。
每个阶段都会基于前一个阶段的权重继续训练。
相比于从零开始单阶段训练，多阶段训练通过逐步扩展数据，更高效地达成高质量视频生成的目标。
△训练方案三阶段
第一阶段是大规模图像预训练。
团队利用互联网上丰富的图像数据和文生图技术，先训练出一个高质量的文生图模型，将该模型作为下一阶段视频预训练的初始化权重。
同时，由于目前没有高质量的时空VAE，他们采用Stable Diffusion预训练好的图像VAE。
这样不仅能保障初始模型的优越性能，还能显著降低视频预训练的整体成本。
第二阶段是大规模视频预训练。
这一阶段主要增加模型的泛化能力，有效掌握视频的时间序列关联。
它需要使用大量视频数据训练，并且保障视频素材的多样性。
同时，第二阶段的模型在第一阶段文生图模型的基础上加入了时序注意力模块，用于学习视频中的时序关系。其余模块与第一阶段保持一致，并加载第一阶段权重作为初始化，同时初始化时序注意力模块输出为零，以达到更高效更快速的收敛。
Colossal-AI团队使用了PixArt-alpha的开源权重作为第二阶段STDiT模型的初始化，以及采用了T5模型作为文本编码器。他们采用了256x256的小分辨率进行预训练，进一步增加了收敛速度，降低训练成本。
△Open-Sora生成效果(提示词：水中世界的镜头，镜头中一只海龟在珊瑚礁间悠然游弋)
第三阶段是高质量视频数据微调。
据介绍，这一阶段能显著提升模型的生成质量。使用的数据规模比上一阶段降低一个量级，但是视频的时长、分辨率和质量都更高。
通过这种方式进行微调，能实现视频生成从短到长、从低分辨率到高分辨率、从低保真度到高保真度的高效扩展。
值得一提的是，Colossal-AI还详细透露了每阶段的资源使用情况。
在Open-Sora的复现流程中，他们使用了64块H800进行训练。第二阶段的训练量一共是 2808 GPU hours，约合7000美元，第三阶段的训练量是1920 GPU hours，大约4500美元。经过初步估算，整个训练方案成功把Open-Sora复现流程控制在了1万美元左右。
数据预处理
为了进一步降低Sora复现的门槛和复杂度，Colossal-AI团队在代码仓库中还提供了便捷的视频数据预处理脚本，让大家可以轻松启动Sora复现预训练。
包括公开视频数据集下载、长视频根据镜头连续性分割为短视频片段、使用开源大语言模型LLaVA生成精细的提示词。
他们提供的批量视频标题生成代码可以用两卡 3 秒标注一个视频，并且质量接近于 GPT-4V。
最终得到的视频/文本对可直接用于训练。借助他们在GitHub上提供的开源代码，可以轻松地在自己的数据集上快速生成训练所需的视频/文本对，显著降低了启动Sora复现项目的技术门槛和前期准备。
高效训练加持
除此之外，Colossal-AI团队还提供了训练加速方案。
通过算子优化和混合并行等高效训练策略，在处理64帧、512x512分辨率视频的训练中，实现了1.55倍的加速效果。
同时，得益于Colossal-AI的异构内存管理系统，在单台服务器上（8H800）可以无阻碍地进行1分钟的1080p高清视频训练任务。
而且团队还发现STDiT模型架构在训练时也展现出卓越的高效性。
和采用全注意力机制的DiT相比，随着帧数的增加，STDiT实现了高达5倍的加速效果，这在处理长视频序列等现实任务中尤为关键。
最后，团队还放出了更多Open-Sora的生成效果。
，时长00:25
团队和量子位透露，他们将长期更新优化Open-Sora的相关解决方案和动态。未来将使用更多视频训练数据，以生成更高质量、更长时长的视频内容，并支持多分辨率特性。
实际应用方面，团队透露将推进在电影、游戏、广告等领域落地。
感兴趣的开发者们，可访问GitHub项目了解更多~
Open-Sora 开源地址：https://github.com/hpcaitech/Open-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。
[7]https://github.com/haotian-liu/LLaVA。
[8]https://hpc-ai.com/blog/open-sora-v1.0。
