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 DatasetsYou 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”文件。
如需更多命令行选项,可以使用以下命令查看帮助信息。
python sample.py -h
请确保视频压缩器和模型架构与训练设置一致。
项目地址:
GitHub - hpcaitech/Open-Sora: Building your own video generation model like OpenAI's Sora
参看文献
OpenAI Sora 技术报告
视频GPT项目
扩散变压器
深速尤利西斯
开放式数字技术
sora相关文章集
好莱坞泰勒佩里Sora制作电影,对好莱坞担忧,停止工作室扩张
可让照片人物“开口说话”阿里图生视频模型EMO,高启强普法
Sora - 探索AI视频模型的无限可能
Sora对于普通人意味着什么?
Sora 在线应用
#2024 年 AI 辅助研发趋势#