Open-Sora,高效复现类Sora视频生成方案开源!魔搭社区最佳实践教程来啦!https://mp.weixin.qq.com/s/WMQIDgZs2MBPGtx18XSXgw
Open-Sora开源方案讲解开源但“平替”的方案。https://mp.weixin.qq.com/s/nPYCzgBA7hIsPZ6PCyXxKQOpen-Sora/docs/report_v1.md at main · hpcaitech/Open-Sora · GitHubOpen-Sora: Democratizing Efficient Video Production for All - Open-Sora/docs/report_v1.md at main · hpcaitech/Open-Sorahttps://github.com/hpcaitech/Open-Sora/blob/main/docs/report_v1.md目前复现sora有两个主力项目,一个是open-sora,另一个是open-sora-plan,前者是hpcaitech,就是colossalAI团队,后者是北大兔展联合实验室推的,这两个都是可以持续跟进的项目,目前简单介绍下open-sora的一些做法。
1.变分自编码器VAE
为了降低计算成本,Open-Sora使用VAE将视频从原始像素空间映射至潜在空间(latent space)。Sora的技术报告中,采用了时空VAE来减少时间维度,Open-Sora项目组通过研究和实践,发现目前尚无开源的高品质时空VAE(3D-VAE)模型。Google研究项目的MAGVIT所使用的4x4x4 VAE并未开放源代码,而VideoGPT的2x4x4 VAE在实验中表现出较低的质量。因此,在Open-Sora v1版本中,使用来自Stability-AI的2D VAE(sd-vae-ft-mse)。
2.Diffusion transformers - STDiT
在处理视频训练时,涉及到大量token。对于每分钟24帧的视频,共有1440帧。经过VAE 4倍下采样和patch尺寸2倍下采样后,大约得到1440x1024≈150万token。对这150万个token进行全注意力操作会导致巨大的计算开销。因此,Open-Sora项目借鉴Latte项目,采用时空注意力机制来降低成本。
如图所示,在STDiT(空间-时间)架构中每个空间注意力模块之后插入一个时间注意力模块。这一设计与Latte论文中的变体3相似,但在参数数量上未做严格控制。Open-Sora在16x256x256分辨率视频上的实验表明,在相同迭代次数下,性能排序为:DiT(全注意力)> STDiT(顺序执行)> STDiT(并行执行)≈ Latte。出于效率考虑,Open-Sora本次选择了STDiT(顺序执行)。
Open-Sora专注视频生成任务,在PixArt-α基础上训练模型。这项研究,采用了T5-conditioned DiT结构。Open-Sora以PixArt-α为基础初始化模型,并将插入的时间注意力层初始化为零值。这样的初始化方式确保了模型从一开始就能保持图像生成能力,文本编码器采用的则是T5模型。插入的时间注意力层使得参数量从5.8亿增加到了7.24亿。
受PixArt-α和稳定视频扩散技术成功的启发,Open-Sora采取了逐步训练策略:首先在36.6万预训练数据集上以16x256x256分辨率训练,然后在2万数据集上分别以16x256x256、16x512x512以及64x512x512分辨率继续训练。借助缩放position embedding,显著降低了计算成本。
3.Patch embedding
Open-Sora尝试在DiT中使用三维patch embedding,但由于在时间维度上进行2倍下采样,生成的视频质量较低。因此,在下一版本中,Open-Sora把下采样的任务留给时空VAE,在V1中按照每3帧采样(16帧训练)和每2帧采样(64帧训练)的方式进行训练。
4.Video caption
Open-Sora使用LLaVA-1.6-Yi-34B为视频进行标注,该标注基于三个连续的帧以及一个精心设计的提示语。借助这个精心设计的提示语,LLaVA能够生成高质量的视频描述。
5.模型架构设计
整个架构包括一个预训练好的VAE,一个文本编码器和一个利用空间-时间注意力机制的STDiT模型,在训练阶段首先采用预训练好的VAE的编码器将视频数据进行压缩,然后在压缩之后的潜在空间中与文本embedding一起训练STDiT,在推理阶段,从VAE的潜在空间中随机采样出一个高斯噪声与提示词prompt embedding一起输入STDiT中,得到去噪之后的特征,最后输入到VAE的解码器中,解码得到视频。
6.训练
训练参考了stable video diffusion,包括三个阶段:1.大规模图像预训练,2.大规模视频预训练;3.高质量视频数据微调。每个阶段都会基于前一个阶段的权重继续训练,相比较从零开始单阶段训练,多阶段训练通过逐步扩展数据,高效。
6.1 第一阶段
通过互联网上丰富的大规模图像数据和先进的文生图技术,训练一个高质量的文生图模型,该模型将作为下一阶段视频预训练的初始化权重,PIxel-alpha的权重。同时,由于目前没有高质量的时空 VAE,采用了 Stable Diffusion [5] 模型预训练好的图像 VAE。该策略不仅保障了初始模型的优越性能,还显著降低了视频预训练的整体成本。
6.2 第二阶段
第二阶段的模型在第一阶段文生图模型的基础上加入了时序注意力模块,用于学习视频中的时序关系。其余模块与第一阶段保持一致,并加载第一阶段权重作为初始化,同时初始化时序注意力模块输出为零,以达到更高效更快速的收敛。采用了T5模型作为文本编码器。同时他们采用了 256x256 的小分辨率进行预训练,进一步增加了收敛速度,降低训练成本。
6.3 第三阶段
第三阶段用到的视频数据规模比第二阶段要少一个量级,但是视频的时长、分辨率和质量都更高。通过这种方式进行微调,他们实现了视频生成从短到长、从低分辨率到高分辨率、从低保真度到高保真度的高效扩展。
Open-Sora 的复现流程中,使用了 64 块 H800 进行训练。第二阶段的训练量一共是 2808 GPU hours,约合 7000 美元,第三阶段的训练量是 1920 GPU hours,大约 4500 美元。经过初步估算,整个训练方案成功把 Open-Sora 复现流程控制在了 1 万美元左右。
7.高效训练
总结来说,Open-Sora项目V1版本非常完整的复刻了基于Transformers的视频生成的Pipeline:
-
比较多种STDiT的方式,采用了STDiT(顺序执行),并验证了结果
-
借助position embedding,实现了不同分辨率和不同时长的视频生成。
在试验的过程中也遇到了一些困难,比如我们注意到,
-
Open-Sora项目一开始采用的是VideoGPT的时空VAE,验证效果不佳后,依然选择了Stable Diffusion的2D的VAE。
-
同时,三维patch embedding,由于在时间维度上进行2倍下采样,生成的视频质量较低。在V1版本中依然采用了按帧采样的方式。
Open-Sora也借助了开源项目和模型的力量,包括但不限于:
-
LLaVA-1.6-Yi-34B的多模态LLM来实现Video-Caption,生成高质量的视频文本对。
-
受PixArt-α和稳定视频扩散技术成功的启发,采用了T5 conditioned DiT结构。