完整复现Sora，Open-Sora最新技术报告发布，提供完整训练代码、权重及数据处理工具 精华

文章链接：https://arxiv.org/pdf/2412.20404
项目链接：https://github.com/hpcaitech/Open-Sora

总结速览

解决的问题
人工视觉智能，特别是生成和模拟我们所见世界的能力，相较于语言能力的突破仍然滞后。现有视频生成模型在高保真视频内容生成、灵活视频合成及长时间视频生成方面面临诸多挑战。

提出的方案
本文引入 Open-Sora，一个开源的视频生成模型，支持文本生成图像、文本生成视频以及图像生成视频等多种视觉生成任务。通过空间-时间扩散Transformer (Spatial-Temporal Diffusion Transformer, STDiT) 框架，将空间与时间的注意力机制解耦，同时采用高度压缩的3D自编码器以压缩表示，加速训练过程。此外，提供完整的训练代码、模型权重及数据处理工具，推动社区发展。

应用的技术

• STDiT框架：高效的扩散视频生成框架，解耦空间和时间注意力。
• 3D自编码器：实现表示的高度压缩，加速训练。
• 定制化训练策略：优化生成效率和效果。

达到的效果

• 支持生成最长15秒、分辨率最高720p的视频，并适配任意宽高比。
• 在文本生成视频、图像生成视频任务中实现可控的运动动态生成。
• 开源以来取得显著成果，模型版本持续更新（v1.0到v1.2），当前最新版本（v1.2）已实现完整复现 OpenAI Sora 的技术，并支持生成多分辨率视频。

数据

数据来源

所使用的数据集全部开源，以确保模型训练的完全可复现性。总计生成了 30M 个视频片段，时长从 2秒到16秒 不等，总时长达 80k小时。

• Webvid-10M：包含10M个来自库存视频网站的视频-文本对。视频为低分辨率并带有水印。
• Panda-70M：一个大规模数据集，包含70M个视频-字幕对。使用了其中20M高质量子集进行训练。
• HD-VG-130M：由130M个文本-视频对组成，字幕通过BLIP-2生成。发现其场景和文本质量相对较差。
• MiraData：一个高质量数据集，包含77k个长视频，主要来源于游戏和城市探索。
• Vript：一个密集标注的数据集，包含400k个视频。
• Inter4K：一个包含1k个4K分辨率视频片段的数据集。

此外，还从 Pexels、Pixabay 和 Mixkit 获取了免费授权的视频。这些网站上的大部分视频质量较高，对这些优秀平台及其贡献者表示由衷的感谢。

图像数据集与视频一起训练，总计包含约 3M 张图像：

• LAION：一个大规模开放数据集，使用了美学评分大于6.5的子集。
• Unsplash-lite：包含25k张自然主题的Unsplash照片，覆盖了广泛的使用场景和上下文。

数据预处理

高质量数据对于训练优秀的生成模型至关重要。为此建立了一条完整的数据处理pipeline，可将原始视频无缝转换为高质量的视频-文本对。pipeline如图2所示。

数据处理步骤：

1. 场景检测与视频剪辑: 使用PySceneCut检测场景并将视频剪辑为多个片段。
2. 高质量视频过滤: 主要遵循SVD 数据预处理pipeline

• 美学评分 (Aesthetic Score)：衡量视频帧的美学偏好。使用来自LAION的评分器，并对每段视频抽样的三帧计算平均分。
• 光流评分 (Optical Flow Score)：衡量视频的动态变化程度，用于筛选低运动量的视频。此任务使用UniMatch模型完成。
• 文本场景过滤：部分视频包含大量密集文字场景（如新闻广播、广告），不适合训练。使用OCR（光学字符识别）检测视频中的文本，含有过多文字的视频将被移除。OCR 使用DBNet++模型，由MMOCR实现。

3. 视频字幕生成: 为视频提供高质量的字幕

• 使用GPT-4V 和 PLLaVA 生成字幕。前者通过 API 提供服务，后者为开源模型，可在本地部署。尽管生成结果存在一定程度的虚构现象，但足以满足训练文本生成视频模型的需求。
• 在实践中，使用预训练的PLLaVA 13B 模型，并从每段视频中选取 4帧 进行字幕生成，采用 2×2 空间池化 形状。
• 字幕生成模型难以提供关于相机运动的信息，通过光流检测相机运动并将其信息附加到字幕中。

最后阶段所用视频数据的统计信息如下图4所示。展示了视频时长和分辨率的基本统计数据，以及美学评分和光流评分的分布。此外，还从视频字幕中提取了与对象和动作相关的标签，并统计了其频率。

模型架构

本文的视频生成框架遵循 Sora 报告 的设计。视频首先通过视频压缩网络（即 3D 自编码器）进行压缩，文本通过文本编码器编码。然后，一个类似 DiT 的 Transformer 处理视频和文本的隐空间变量。

3D 自编码器

在 Open-Sora 1.0 和 1.1 中，使用了 Stability-AI 的 2D VAE（84M 参数），该模型以 8×8 的比例对视频进行空间压缩。为减少时间维度，通过每三帧提取一帧进行下采样。然而，这种方法由于生成的帧率（FPS）降低，导致时间流畅性较差。

为解决这一限制，Open-Sora 1.2 引入了一个受 OpenAI 的 Sora 启发的视频压缩网络，在时间维度上实现了 4 倍压缩。这消除了帧提取的需求，使视频能够以原始帧率生成。

由于训练 3D VAE 的计算需求较高，尝试利用 2D VAE 中嵌入的知识。通过 2D VAE 压缩后，观察到时间上相邻的特征高度相关。基于这一观察，开发了一个简单但有效的视频压缩网络，先以 8×8 的比例进行空间压缩，然后在时间上以 4 倍压缩。网络架构如下图5所示。

• 2D VAE 初始化：使用SDXL 的预训练 VAE。
• 3D VAE 架构：采用Magvit-v2 的 VAE结构，总参数量为300M。结合2D VAE后，视频压缩网络的总参数量为384M。
• 训练细节：3D VAE 使用来自Pexels和Pixabay的视频进行训练，总共训练了1.2M 步，本地批量大小为1。训练数据主要为分辨率256×256的17 帧短视频片段。
• 优化：在 3D VAE 中使用了因果卷积以提高图像重建精度。

训练过程
训练过程分为三个阶段：

• 阶段 1（0–380k 步）：使用8 张 GPU进行训练，冻结2D VAE的权重。目标包括重建2D VAE压缩的特征，并应用身份损失（identity loss）来对齐3D VAE和2D VAE的特征。身份损失加速了收敛速度，并提高了初始图像重建质量。
• 阶段 2（380k–640k 步）：移除了身份损失，3D VAE 通过训练进一步优化其时间维度的理解能力。
• 阶段 3（640k–1.2M 步）：发现仅重建2D VAE特征不足以进一步提升模型性能，因此损失函数被替换为直接重建原始视频。在此阶段，使用24 张 GPU，并通过随机化视频长度（最长34 帧）结合适当的零填充实现混合视频长度训练，从而提高了模型对不同视频时长的鲁棒性。

在前两个阶段，数据集由 80% 视频数据 和 20% 图像数据 组成。对于视频训练，使用 17 帧视频片段；图像数据通过零填充与视频输入格式匹配。然而，这种方法导致了非标准长度视频的模糊问题，而第三阶段的混合长度训练有效解决了这一问题。

堆叠的 VAE 架构在推理时占用的内存极少，因为输入已经经过压缩。为了提高效率，输入视频被分割为 17 帧的片段。与另一个开源 3D VAE 相比，本文的模型在显著降低计算成本的同时，性能与其相当。

架构

本文的模型架构基于 PixArt，这是一种图像扩散 Transformer。文本通过 T5 文本编码器 进行编码，视频和文本隐空间变量之间采用 交叉注意力 机制。

为了实现高效的视频生成，本文引入了一种空间-时间注意力机制，即 空间-时间扩散 Transformer（STDiT），灵感来源于 Latte，替代了对所有tokens的完全注意力机制。具体来说：

• 空间自注意力在每一帧内应用；
• 时间注意力在相同空间位置的帧之间应用。

为了专注于视频生成，设计了基于强大的预训练图像生成模型的架构。模型以 PixArt-α 初始化，这是一种 T5 条件 DiT 结构，针对高质量和高效图像生成进行了优化。为新引入的时间注意力设计的投影层初始值设为 0，以保持模型在训练开始时的原始图像生成能力。

时间注意力的加入使参数数量从 580M 增加到 1.1B。

条件控制

尽管文本到视频生成具有高度的多样性，但某些应用场景需要更精确的控制。为了实现这一点，通过引入图像和视频输入的遮罩策略，扩展了模型的功能，支持图像到图像和视频到视频生成。如下图 7 所示。

遮罩策略
在该方法中：

• 被指定为条件控制的帧会被取消遮罩。
• 在前向传播过程中：

• 取消遮罩的帧被赋予时间步为 0。
• 其他帧保持其扩散时间步。

然而，将该策略直接应用于预训练模型通常会产生次优结果，因为扩散模型未经过处理单个样本中混合时间步的训练。

受到 UL2 的启发，通过随机遮罩策略解决了这一问题。具体来说，帧会以随机模式取消遮罩，例如：

• 第一帧
• 前k帧
• 最后一帧
• 后k帧
• 前后k帧的组合
• 或完全随机的帧。

在 Open-Sora 1.0 的基础上，对 50% 的训练样本应用mask策略进行实验，发现模型在 10k 步之后能够有效学习图像和视频的条件控制能力，同时对文本到视频性能的影响最小。而较低的mask概率（例如 30%）会降低条件控制的效果。因此，从头开始使用此mask策略预训练模型。

为了进一步增强模型的控制能力，在字幕中附加分数，并将其作为附加条件输入。这些分数包括美学分数、运动分数和相机运动描述。例如，一个美学分数为 5.5、运动分数为 10，并检测到相机运动为“向左平移”的视频，其字幕格式为：
[原始字幕] 美学分数: 5.5，运动分数: 10，相机运动: 向左平移。

在推理过程中，这些分数也可以调整以影响视频生成。对于相机运动条件控制，手动标注了 13,000 个高置信度的片段。

这种方法使模型能够对条件输入有更细致的理解，从而提高了其在各种任务中生成高质量、上下文相关视频的能力。

训练策略

多分辨率和多长宽比

正如 Sora 的报告中所指出的，使用视频的原始分辨率、长宽比和长度进行训练能够提高采样的灵活性，并增强画面构图能力。为实现这一目标，评估了三种方法：

1. NaViT

• 通过遮罩支持同一批次内的动态尺寸，效率损失最小。
• 实现复杂，可能无法充分利用如 Flash Attention [8] 之类的优化内核。

2. 填充（FiT）

• 通过将较小分辨率填充到最大尺寸，支持同一批次内的动态尺寸。
• 实现简单，但对分辨率变化较大的情况下，内存使用效率较低。

3. 分桶（Bucket）（SDXL, PixArt）

• 通过将样本分组到预定义的“桶”中支持跨批次的动态尺寸。
• 每个批次内分辨率、帧数和长宽比固定。分桶方法避免了遮罩或填充的复杂性，同时在统一尺寸的输入上可以利用优化操作。
• 然而，这种方法的灵活性受到预定义尺寸集合的限制。

为简单高效起见，采用了基于分桶的方法。预定义了一组固定的分辨率、长宽比和帧长度，并据此将样本分配到相应的桶中。每个桶由三元组 分辨率帧数长宽比 定义，以覆盖大多数常见的视频格式。在每个训练周期前，数据集会被重新打乱，样本被分配到能够容纳其分辨率和帧长度的最大桶中。

为了进一步优化计算资源，为每个桶引入了两个附加属性：

• 保留概率（Probability of Keeping in the Bucket）
  高分辨率视频会根据概率被下采样为较低分辨率，从而有效降低计算成本。
• 批次大小（Batch Size）
  根据每个桶的情况调整批次大小，以平衡 GPU 负载，确保资源利用效率。

通过对这些参数进行微调，我们在桶之间实现了样本的均衡分布，在保持高质量视频生成的同时提升了总体训练效率。

这种基于分桶的策略在实现简单性与计算效率之间提供了一个实用的折中，支持具有多样化分辨率和长宽比的视频训练。

模型适配

从 PixArt-Σ 2K checkpoint 开始，该模型使用 DDPM 和 SDXL VAE，在更高分辨率下进行训练。通过在较小数据集上微调，该模型被高效地适配到视频生成任务中。适配过程包括多个连续阶段，全部在 8 个 H100 GPU 上完成：

1. 多分辨率图像生成

• 训练模型以支持从 144p 到 2K 的分辨率，共 20k 步。

2. 引入 QK 正则化

• 为提高稳定性，添加 QK-norm，训练 18k 步。

3. 转向修正流

• 从离散时间的 DDPM 过渡到连续时间的修正流，共 10k 步。

4. 强化修正流训练

• 加入 logit-norm 采样和分辨率感知时间步采样，训练 33k 步。

5. 更小的 AdamW epsilon

6. 新 VAE 和 FPS 条件控制

• 用 Open-Sora 的 VAE 替换原有 VAE，添加 FPS 条件控制到时间步条件控制中，训练 25k 步。
• 在修正流训练中，对每个通道进行归一化被证明是关键步骤。

7. 时间注意力模块

• 添加零初始化的时间注意力模块，初始仅在图像上训练 3k 步。

8. 时间模块的mask策略

• 使用mask策略专注于视频的时间注意力模块，训练 38k 步。

适配后模型的优势

完成适配后，模型不仅保留了生成高质量图像的能力，还在视频生成方面获得了多重优势：

1. 加速训练和推理

• 修正流将视频采样步骤从 100 减少到 30，显著降低推理时间。

2. 增强稳定性

• QK-norm 支持更激进的优化，提升训练效率。

3. 高效的时间维度压缩

• 新 VAE 将时间维度压缩至原来的四分之一，减少计算成本。

4. 分辨率灵活性

• 模型能够生成多种分辨率的视频，从 144p 到 2K，支持多样化应用场景。

这种全面的适配不仅提升了模型的视频生成能力，还确保了训练的高效性和可扩展性，为开源基于扩散的视频生成树立了新的标准。

多阶段训练

为了在有限的计算预算内优化性能，我们根据数据质量精心组织训练数据，并将训练过程分为三个阶段。模型在一个 12×8 GPU 的设置上进行了大约两周的训练，完成了约 70k 步。

在第一阶段，模型在 Webvid-10M 数据集（包含 40,000 小时的视频）上进行了 30k 步训练（2 个 epoch）。该数据集主要包含分辨率低于 360p 且带有水印的视频，非常适合初步训练。我们专注于 240p 和 360p 分辨率的视频，长度为 2 到 16 秒。训练使用了原始数据集的字幕。

在第二阶段，使用 Panda-70M 数据集 进行了训练。由于该数据集的质量不均，使用了官方的 30M 子集，并筛选出只有美学分数高于 4.5 的视频，最终得到 20M 的子集（41,000 小时）。训练主要集中在 360p 和 480p 分辨率的视频上，共进行了 23k 步训练，约为 0.5 个 epoch。虽然这一阶段的训练没有完全完成，但它为模型的广泛应用提供了足够的改进。

最后阶段包括从各种来源精选的约 2M 高质量视频片段，总计 5,000 小时。这些视频来自 MiraData 和 Vript，并由 GPT 生成了字幕，其他来源则使用 PLLaVA 标注。该阶段主要集中在更高分辨率（720p 和 1080p）的训练，以提升模型处理更大分辨率的能力。在训练过程中，采用了 25% 的mask比率，训练共进行了 15k 步（约 2 个 epoch）。

为了进行验证，从 Pixabay 中随机抽取了 1k 个视频来评估模型的性能。评估损失值针对不同长度（2s、4s、8s、16s）和不同分辨率（144p、240p、360p、480p、720p）的图像和视频进行了计算。每个配置的损失值在 10 个等距的时间步上进行了平均。

还在训练过程中跟踪了 VBench 分数。VBench 是一个用于评估短视频生成的自动化基准工具。使用 240p 2 秒的视频计算分数，为模型的进展提供了额外的验证。评估损失和 VBench 分数都确认了模型在整个训练过程中不断取得进展。训练过程中的 VBench 分数和验证损失见下图 9。

下表 2 显示了各个模型的 VBench 分数，证明了 Open-Sora 在开源视频生成模型中达到了行业领先水平。

结论

Open-Sora 代表了开源视频生成领域的一次重大进步，提供了一个全面的框架，包括数据处理、训练代码和模型权重。通过成功地重现了 Sora 报告中的关键技术，并实现了最长 16 秒的视频生成，分辨率可达 720p，且支持可控的运动动态，Open-Sora 使得先进的视频生成技术得以普及。该项目不仅促进了社区的合作，还为该领域未来的发展奠定了基础。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/2JnePAI3YzBWrJj1ExlKWg​​