Open-Sora 开源版Sora复现方案

🎥 前言

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

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

📎 目录

•🎥 Open-Sora     •📎 目录     •📍 概览     •📂 数据集准备          •使用 MSR-VTT          •使用自定义数据集      •🚀 开始使用          •训练          •推理

最新消息

[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/MSRVTTbash 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/trainpython scripts/data/preprocess_data.py –c ./dataset/MSRVTT–collated/val/annotations.json –v ./dataset/MSRVTT–collated/val/videos –o ./dataset/MSRVTT–processed/valpython 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