近日,Colossal-AI团队开源了Open-Sora 1.0——一款基于Diffusion Transformer(DiT)架构的类Sora视频生成模型,其核心组件STDiT(Spatial Temporal Diffusion Transformer)融合了空间-时间注意力机制,创新性地将已有的高质量文生图模型PixArt-α拓展至视频领域。STDiT结构巧妙地串联起二维空间注意力模块和一维时间注意力模块,精准模拟视频数据的时序关联,而在时间注意力模块后紧跟的交叉注意力集中模块则确保了文本语义与生成视频的深度对齐。这种结构相较于全注意力机制显著减少了计算资源需求,同时也优于同类空间-时间注意力模型Latte,能够更好地利用预训练好的图像DiT权重迁移学习至视频场景。
三阶段训练复现方案详析
第一阶段:大规模图像预训练降低成本

Open-Sora复现过程中,首要阶段即采用大规模图像数据预训练成熟文生图模型,以此来初始化后续的视频预训练,从而大幅度削减成本。Colossal-AI团队借鉴Stable Diffusion模型的预训练图像VAE,既保证了初始模型性能,又极大节省了视频预训练投入。
第二阶段:大规模视频预训练增强泛化
第二阶段转向大规模视频预训练,通过引入时序注意力模块,使模型能理解视频中的时间序列关系。Colossal-AI团队利用PixArt-alpha开源权重初始化STDiT模型,并结合T5作为文本编码器,以256x256低分辨率预训练加快收敛速度,提高训练效率。
第三阶段:高质量视频微调提升精度
在第三阶段,团队对高质量视频数据进行了精细化微调,虽然数据量减少,但视频长度、分辨率及画质均有所提升,实现了从短片、低清到长片、高清的高效过渡。这一阶段微调使得Open-Sora生成的视频质量明显跃升。
数据预处理自动化简化流程
为便于用户快速启动Sora复现预训练,Colossal-AI贴心提供了易用的视频数据预处理脚本。该脚本囊括了公开视频数据集下载、长视频智能切割、以及运用LLaVA开源大语言模型自动生成精细提示词等功能。团队开发的批量视频标题生成工具高效快捷,每3秒即可完成一个视频的标注,生成质量堪比GPT-4V水平,由此产出的视频/文本对可以直接投入训练。
视频生成效果演示与展望
Open-Sora的实际应用案例令人惊叹不已。模型能够依据指令创造出如悬崖海岸航拍、山川瀑布奔腾入湖、海底世界海龟漫游、星空银河延时摄影等各种生动场景。尽管当前版本仅使用了400K训练数据,存在生成质量及文本贴合度有待提高的问题(如生成的海龟多出一只脚),但团队已在GitHub上明确列出改进计划,矢志不渝地优化模型性能。
高效训练技术的实践与突破
除了降低复现技术门槛和提升视频生成质量,Colossal-AI团队还通过Colossal-AI加速系统,采用算子优化和混合并行策略,实现对64帧、512x512分辨率视频训练高达1.55倍的加速。此外,得益于异构内存管理系统,在配备8个H800显卡的服务器上,能够流畅进行1分钟1080p高清视频训练。
值得一提的是,STDiT模型在训练中表现出非凡的高效性,特别是在处理长视频序列时,其加速效果远超全注意力机制的DiT,达到5倍以上的提升。
Open-Sora项目的开源标志着视频生成技术向前迈出了重要一步,Colossal-AI团队将持续致力于优化Open-Sora,增加更多视频训练数据、支持更高质量和更长时间的视频生成,并积极探索在影视制作、游戏开发、广告创意等领域内的广泛应用。敬请广大AI爱好者密切关注Open-Sora的开源社区(https://github.com/hpcaitech/Open-Sora),一同见证并参与这一技术的发展历程。
来源:公众号“小微模型”