Open-Sora：开源 Sora 复现方案，成本降低 46%

演示站点： https://ai.uaai.cn 官方论坛： www.jingyuai.com 京娱AI，提供完整的 Sora 复现架构方案，包含从数据处理到训练推理全流程。支持动态分辨率，训练时可直接训练任意分辨率的视频，无需进行缩放。支持多种模型结构。由于 Sora 实际模型结构未知，用户可自行选择使用原始视频、VQVAE（视频原生的模型）、SD-VAE（图像原生的模型）进行训练。

yimifx

1497人浏览 · 2024-03-10 00:12:36

yimifx · 2024-03-10 00:12:36 发布

Colossal-AI 开源了完整的 Sora 复现架构方案 Open-Sora，声称可降低 46% 复现成本，并将模型训练输入序列长度扩充至 819K patches。

演示站点： https://ai.uaai.cn UAAI
官方论坛： www.jingyuai.com京娱AI

Sora 算法复现方案

在 Sora 的技术报告中，Sora 使用了一个视频压缩网络将各种尺寸的视频压缩成一个隐空间 (latent space) 的时空块序列 (a sequence of patial temporal patch)，然后使用了 Diffusion Transformer 进行去噪，最后进行解码生成视频。

Open-Sora 将 Sora 可能使用的训练 pipeline 归纳为下图。

📍 概述

目前 Open-Sora 已涵盖：

提供完整的 Sora 复现架构方案，包含从数据处理到训练推理全流程。
支持动态分辨率，训练时可直接训练任意分辨率的视频，无需进行缩放。
支持多种模型结构。由于 Sora 实际模型结构未知，我们实现了 adaLN-zero、cross attention、in-context conditioning (token concat) 等三种常见的多模态模型结构。
支持多种视频压缩方法。用户可自行选择使用原始视频、VQVAE（视频原生的模型）、SD-VAE（图像原生的模型）进行训练。
支持多种并行训练优化。包括结合 Colossal-AI 的 AI 大模型系统优化能力，及 Ulysses 和 FastSeq 的混合序列并行。

性能

以在单台 H800 SXM 8*80GB GPU 上使用 DiT-XL/2 模型的性能测试为例。在 600K 的序列长度时，Open-Sora 的方案比基线方案有 40% 以上的性能提升和成本降低。

📂 数据集准备

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

# Step 1: download the dataset to ./dataset/MSRVTT
bash scripts/data/download_msr_vtt_dataset.sh

# Step 2: collate the video and annotations
python scripts/data/collate_msr_vtt_dataset.py -d ./dataset/MSRVTT/ -o ./dataset/MSRVTT-collated

# Step 3: perform data processing
# NOTE: each script could several minutes so we apply the script to the dataset split individually
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，你应该--use_pooled_text在运行时使用preprocess_data.py


After completing these steps, you should have a processed MSR-VTT dataset in `./dataset/MSRVTT-processed`.


### Use Customized Datasets

You can also use other datasets and transform the dataset to the required format. You should prepare a captions file and a video directory. The captions file should be a JSON file or a JSONL file. The video directory should contain all the videos.

Here is an example of the captions file:

```json
[
    {
        "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来量化视频帧。我们使用CLIP来提取文本特征。

输出是一个箭头数据集，其中包含以下列：“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，你应该--use_pooled_text在运行时使用preprocess_data.py

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

🚀 开始吧

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

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”文件。

如需更多命令行选项，可以使用以下命令查看帮助信息。