Open-Sora 开源版Sora复现方案

🎥 前言

在人工智能盛起的当下，AI正以非常迅猛的速度重塑着很多行业。可以预见的是2024将是AI原生应用开发元年，将会涌现出数不清的AI原生应用来重塑我们的工作和生活的方方面面。而在AI原生应用里面将会以AI Agent即AI智能体为主要代表，将会有很多个像crewAI—用于编排角色扮演的AI agent（超级智能体）一样的Agent出现在我们的面前。在可以预见的未来，世界大模型Sora—聊聊火出圈的世界AI大模型——Sora毫无疑问将会带来革命性的AI热潮。

今天要介绍的是一款复现Sora技术的开源项目。

最新消息

[2024/03] Open-Sora：Sora 复现解决方案，成本降低 46%，序列扩展至近百万[1]

📍 概览

Open-Sora 是一个开源项目，提供了一个高性能的开发管道实现，Sora 可能使用这种管道，由 Colossal-AI[2]驱动，包括：

•提供完整的 Sora 复现架构解决方案，包括从数据处理到训练和部署的整个过程。

•支持动态分辨率，训练可以直接训练任何分辨率的视频，无需缩放。

•支持多种模型结构。由于 Sora 的实际模型结构未知，我们实现了三种常见的多模态模型结构，如 adaLN-zero、交叉注意力和上下文条件化（token concat）。

•支持多种视频压缩方法。用户可以选择使用原始视频、VQVAE（视频原生模型）或 SD-VAE（图像原生模型）进行训练。

•支持多种并行训练优化。包括与 Colossal-AI 兼容的 AI 大模型系统优化能力，以及 Ulysses 和 FastSeq 的混合序列并行性。

📂 数据集准备

使用 MSR-VTT

我们使用 MSR-VTT[3] 数据集，这是一个大规模的视频描述数据集。用户应在训练模型之前对原始视频进行预处理。您可以使用以下脚本进行数据处理。

# 第一步：下载数据集到 ./dataset/MSRVTT
bash scripts/data/download_msr_vtt_dataset.sh

# 第二步：整理视频和注释
python scripts/data/collate_msr_vtt_dataset.py -d ./dataset/MSRVTT/ -o ./dataset/MSRVTT-collated

# 第三步：进行数据处理
# 注意：每个脚本可能需要几分钟，因此我们分别对数据集划分应用脚本
python scripts/data/preprocess_data.py -c ./dataset/MSRVTT-collated/train/annotations.json -v ./dataset/MSRVTT-collated/train/videos -o ./dataset/MSRVTT-processed/train
python scripts/data/preprocess_data.py -c ./dataset/MSRVTT-collated/val/annotations.json -v ./dataset/MSRVTT-collated/val/videos -o ./dataset/MSRVTT-processed/val
python scripts/data/preprocess_data.py -c ./dataset/MSRVTT-collated/test/annotations.json -v ./dataset/MSRVTT-collated/test/videos -o ./dataset/MSRVTT-processed/test

如果您想使用 adaLN-zero，应在运行 preprocess_data.py 时使用 --use_pooled_text

完成这些步骤后，您应该在 ./dataset/MSRVTT-processed 中拥有一个处理过的 MSR-VTT 数据集。

使用自定义数据集
您也可以使用其他数据集，并将数据集转换为所需格式。您应该准备一个字幕文件和一个视频目录。字幕文件应该是一个 JSON 文件或一个 JSONL 文件。视频目录应包含所有视频。

这里是字幕文件的一个示例：

[
    {
        "file": "video0.mp4",
        "captions": ["a girl is throwing away folded clothes", "a girl throwing cloths around"]
    },
    {
        "file": "video1.mp4",
        "captions": ["a comparison of two opposing team football athletes"]
    }
]

这里是视频目录的一个示例：

.
├── video0.mp4
├── video1.mp4
└── ...

每个视频可能有多个字幕。因此，输出是视频-字幕对。例如，如果第一个视频有两个字幕，那么输出将是两个视频-字幕对。

我们使用 VQ-VAE[4] 对视频帧进行量化。我们使用 CLIP[5] 提取文本特征。

输出是一个 arrow 数据集，包含以下列："video_file"、"video_latent_states"、"text_latent_states"。"video_latent_states" 的维度是 (T, H, W)，"text_latent_states" 的维度是 (S, D)。

然后，您可以使用以下命令运行数据处理脚本：

python preprocess_data.py -c /path/to/captions.json -v /path/to/video_dir -o /path/to/output_dir

如果您想使用 adaLN-zero，应在运行 preprocess_data.py 时使用--use_pooled_text

请注意，这个脚本需要在带有 GPU 的机器上运行。为了避免 CUDA OOM，我们过滤掉了太长的视频。

🚀 开始使用

在本节中，我们将指导如何运行训练和推理。在此之前，请确保您使用以下命令安装了依赖项。

pip install -r requirements.txt

训练

您可以通过以下命令调用训练。

bash ./scripts/train.sh

您也可以根据需要修改 train.sh 中的参数。

我们提供了三种模型架构："adaln"、"cross-attn" 和 "token-concat"。您可以使用-x 选项指定模型架构。

推理

我们提供了一个脚本来执行推理，允许您从训练好的模型生成视频。您可以通过以下命令调用推理。

python sample.py -m "DiT/XL-2" --text "a person is walking on the street" --ckpt /path/to/checkpoint --height 256 --width 256 --fps 10 --sec 5 --disable-cfg

这将在当前目录生成一个 "sample.mp4" 文件。

要了解更多命令行选项，您可以使用以下命令查看帮助信息。

python sample.py -h

请确保视频压缩器和模型架构与训练设置一致。

其他

本文由山行翻译整理自：https://github.com/hpcaitech/Open-Sora，如对您有帮助，请帮忙关注、点赞、转发。

References

[1] Open-Sora：Sora 复现解决方案，成本降低 46%，序列扩展至近百万: https://hpc-ai.com/blog/open-sora [2] Colossal-AI: https://github.com/hpcaitech/ColossalAI [3] MSR-VTT: https://cove.thecvf.com/datasets/839 [4] VQ-VAE: https://github.com/wilson1yan/VideoGPT/ [5] CLIP: https://huggingface.co/docs/transformers/model_doc/clip#clip