不久前 OpenAI Sora 以其驚人的視頻生成效果迅速走紅，在一眾文生視頻模型中突出重圍，成為全球矚目的焦點(diǎn)。繼 2 周前推出成本直降 46% 的 Sora 訓(xùn)練推理復(fù)現(xiàn)流程后，Colossal-AI 團(tuán)隊(duì)全面開源全球首個(gè)類 Sora 架構(gòu)視頻生成模型 「Open-Sora 1.0」，涵蓋了整個(gè)訓(xùn)練流程，包括數(shù)據(jù)處理、所有訓(xùn)練細(xì)節(jié)和模型權(quán)重，攜手全球 AI 熱愛者共同推進(jìn)視頻創(chuàng)作的新紀(jì)元。

先睹為快，我們先看一段由 Colossal-AI 團(tuán)隊(duì)發(fā)布的「Open-Sora 1.0」模型生成的都市繁華掠影視頻。

Open-Sora 1.0 生成的都市繁華掠影

這僅僅是 Sora 復(fù)現(xiàn)技術(shù)冰山的一角，關(guān)于以上文生視頻的模型架構(gòu)、訓(xùn)練好的模型權(quán)重、復(fù)現(xiàn)的所有訓(xùn)練細(xì)節(jié)、數(shù)據(jù)預(yù)處理過程、demo 展示和詳細(xì)的上手教程，Colossal-AI 團(tuán)隊(duì)已經(jīng)全面免費(fèi)開源在 GitHub，同時(shí)筆者第一時(shí)間聯(lián)系了該團(tuán)隊(duì)，了解到他們將不斷更新 Open-Sora 的相關(guān)解決方案和最新動(dòng)態(tài)，感興趣的朋友可以持續(xù)關(guān)注 Open-Sora 的開源社區(qū)。

Open-Sora 開源地址：https://github.com/hpcaitech/Open-Sora

全面解讀 Sora 復(fù)現(xiàn)方案

接下來，我們將深入解讀 Sora 復(fù)現(xiàn)方案的多個(gè)關(guān)鍵維度，包括模型架構(gòu)設(shè)計(jì)、訓(xùn)練復(fù)現(xiàn)方案、數(shù)據(jù)預(yù)處理、模型生成效果展示以及高效訓(xùn)練優(yōu)化策略。

模型架構(gòu)設(shè)計(jì)

模型采用了目前火熱的 Diffusion Transformer (DiT) [1] 架構(gòu)。作者團(tuán)隊(duì)以同樣使用 DiT 架構(gòu)的高質(zhì)量開源文生圖模型 PixArt-α [2] 為基座，在此基礎(chǔ)上引入時(shí)間注意力層，將其擴(kuò)展到了視頻數(shù)據(jù)上。具體來說，整個(gè)架構(gòu)包括一個(gè)預(yù)訓(xùn)練好的 VAE，一個(gè)文本編碼器，和一個(gè)利用空間 - 時(shí)間注意力機(jī)制的 STDiT (Spatial Temporal Diffusion Transformer) 模型。其中，STDiT 每層的結(jié)構(gòu)如下圖所示。它采用串行的方式在二維的空間注意力模塊上疊加一維的時(shí)間注意力模塊，用于建模時(shí)序關(guān)系。在時(shí)間注意力模塊之后，交叉注意力模塊用于對(duì)齊文本的語意。與全注意力機(jī)制相比，這樣的結(jié)構(gòu)大大降低了訓(xùn)練和推理開銷。與同樣使用空間 - 時(shí)間注意力機(jī)制的 Latte [3] 模型相比，STDiT 可以更好的利用已經(jīng)預(yù)訓(xùn)練好的圖像 DiT 的權(quán)重，從而在視頻數(shù)據(jù)上繼續(xù)訓(xùn)練。

STDiT 結(jié)構(gòu)示意圖

整個(gè)模型的訓(xùn)練和推理流程如下。據(jù)了解，在訓(xùn)練階段首先采用預(yù)訓(xùn)練好的 Variational Autoencoder (VAE) 的編碼器將視頻數(shù)據(jù)進(jìn)行壓縮，然后在壓縮之后的潛在空間中與文本嵌入 (text embedding) 一起訓(xùn)練 STDiT 擴(kuò)散模型。在推理階段，從 VAE 的潛在空間中隨機(jī)采樣出一個(gè)高斯噪聲，與提示詞嵌入 (prompt embedding) 一起輸入到 STDiT 中，得到去噪之后的特征，最后輸入到 VAE 的解碼器，解碼得到視頻。

模型的訓(xùn)練流程

訓(xùn)練復(fù)現(xiàn)方案

我們向該團(tuán)隊(duì)了解到，Open-Sora 的復(fù)現(xiàn)方案參考了 Stable Video Diffusion (SVD)[3] 工作，共包括三個(gè)階段，分別是：

• 大規(guī)模圖像預(yù)訓(xùn)練。
• 大規(guī)模視頻預(yù)訓(xùn)練。
• 高質(zhì)量視頻數(shù)據(jù)微調(diào)。

每個(gè)階段都會(huì)基于前一個(gè)階段的權(quán)重繼續(xù)訓(xùn)練。相比于從零開始單階段訓(xùn)練，多階段訓(xùn)練通過逐步擴(kuò)展數(shù)據(jù)，更高效地達(dá)成高質(zhì)量視頻生成的目標(biāo)。

訓(xùn)練方案三階段

第一階段：大規(guī)模圖像預(yù)訓(xùn)練

第一階段通過大規(guī)模圖像預(yù)訓(xùn)練，借助成熟的文生圖模型，有效降低視頻預(yù)訓(xùn)練成本。

作者團(tuán)隊(duì)向我們透露，通過互聯(lián)網(wǎng)上豐富的大規(guī)模圖像數(shù)據(jù)和先進(jìn)的文生圖技術(shù)，我們可以訓(xùn)練一個(gè)高質(zhì)量的文生圖模型，該模型將作為下一階段視頻預(yù)訓(xùn)練的初始化權(quán)重。同時(shí)，由于目前沒有高質(zhì)量的時(shí)空 VAE，他們采用了 Stable Diffusion [5] 模型預(yù)訓(xùn)練好的圖像 VAE。該策略不僅保障了初始模型的優(yōu)越性能，還顯著降低了視頻預(yù)訓(xùn)練的整體成本。

第二階段：大規(guī)模視頻預(yù)訓(xùn)練

第二階段執(zhí)行大規(guī)模視頻預(yù)訓(xùn)練，增加模型泛化能力，有效掌握視頻的時(shí)間序列關(guān)聯(lián)。

我們了解到，這個(gè)階段需要使用大量視頻數(shù)據(jù)訓(xùn)練，保證視頻題材的多樣性，從而增加模型的泛化能力。第二階段的模型在第一階段文生圖模型的基礎(chǔ)上加入了時(shí)序注意力模塊，用于學(xué)習(xí)視頻中的時(shí)序關(guān)系。其余模塊與第一階段保持一致，并加載第一階段權(quán)重作為初始化，同時(shí)初始化時(shí)序注意力模塊輸出為零，以達(dá)到更高效更快速的收斂。Colossal-AI 團(tuán)隊(duì)使用了 PixArt-alpha [2] 的開源權(quán)重作為第二階段 STDiT 模型的初始化，以及采用了 T5 [6] 模型作為文本編碼器。同時(shí)他們采用了 256x256 的小分辨率進(jìn)行預(yù)訓(xùn)練，進(jìn)一步增加了收斂速度，降低訓(xùn)練成本。

第三階段：高質(zhì)量視頻數(shù)據(jù)微調(diào)

第三階段對(duì)高質(zhì)量視頻數(shù)據(jù)進(jìn)行微調(diào)，顯著提升視頻生成的質(zhì)量。

作者團(tuán)隊(duì)提及第三階段用到的視頻數(shù)據(jù)規(guī)模比第二階段要少一個(gè)量級(jí)，但是視頻的時(shí)長(zhǎng)、分辨率和質(zhì)量都更高。通過這種方式進(jìn)行微調(diào)，他們實(shí)現(xiàn)了視頻生成從短到長(zhǎng)、從低分辨率到高分辨率、從低保真度到高保真度的高效擴(kuò)展。

作者團(tuán)隊(duì)表示，在 Open-Sora 的復(fù)現(xiàn)流程中，他們使用了 64 塊 H800 進(jìn)行訓(xùn)練。第二階段的訓(xùn)練量一共是 2808 GPU hours，約合 7000 美元，第三階段的訓(xùn)練量是 1920 GPU hours，大約 4500 美元。經(jīng)過初步估算，整個(gè)訓(xùn)練方案成功把 Open-Sora 復(fù)現(xiàn)流程控制在了 1 萬美元左右。

數(shù)據(jù)預(yù)處理

為了進(jìn)一步降低 Sora 復(fù)現(xiàn)的門檻和復(fù)雜度，Colossal-AI 團(tuán)隊(duì)在代碼倉庫中還提供了便捷的視頻數(shù)據(jù)預(yù)處理腳本，讓大家可以輕松啟動(dòng) Sora 復(fù)現(xiàn)預(yù)訓(xùn)練，包括公開視頻數(shù)據(jù)集下載，長(zhǎng)視頻根據(jù)鏡頭連續(xù)性分割為短視頻片段，使用開源大語言模型 LLaVA [7] 生成精細(xì)的提示詞。作者團(tuán)隊(duì)提到他們提供的批量視頻標(biāo)題生成代碼可以用兩卡 3 秒標(biāo)注一個(gè)視頻，并且質(zhì)量接近于 GPT-4V。最終得到的視頻 / 文本對(duì)可直接用于訓(xùn)練。借助他們?cè)?GitHub 上提供的開源代碼，我們可以輕松地在自己的數(shù)據(jù)集上快速生成訓(xùn)練所需的視頻 / 文本對(duì)，顯著降低了啟動(dòng) Sora 復(fù)現(xiàn)項(xiàng)目的技術(shù)門檻和前期準(zhǔn)備。

基于數(shù)據(jù)預(yù)處理腳本自動(dòng)生成的視頻 / 文本對(duì)

模型生成效果展示

下面我們來看一下 Open-Sora 實(shí)際視頻生成效果。比如讓 Open-Sora 生成一段在懸崖海岸邊，海水拍打著巖石的航拍畫面。

再讓 Open-Sora 去捕捉山川瀑布從懸崖上澎湃而下，最終匯入湖泊的宏偉鳥瞰畫面。

除了上天還能入海，簡(jiǎn)單輸入 prompt，讓 Open-Sora 生成了一段水中世界的鏡頭，鏡頭中一只海龜在珊瑚礁間悠然游弋。

Open-Sora 還能通過延時(shí)攝影的手法，向我們展現(xiàn)了繁星閃爍的銀河。

如果你還有更多視頻生成的有趣想法，可以訪問 Open-Sora 開源社區(qū)獲取模型權(quán)重進(jìn)行免費(fèi)的體驗(yàn)。鏈接：https://github.com/hpcaitech/Open-Sora

值得注意的是，作者團(tuán)隊(duì)在 Github 上提到目前版本僅使用了 400K 的訓(xùn)練數(shù)據(jù)，模型的生成質(zhì)量和遵循文本的能力都有待提升。例如在上面的烏龜視頻中，生成的烏龜多了一只腳。Open-Sora 1.0 也并不擅長(zhǎng)生成人像和復(fù)雜畫面。作者團(tuán)隊(duì)在 Github 上列舉了一系列待做規(guī)劃，旨在不斷解決現(xiàn)有缺陷，提升生成質(zhì)量。

高效訓(xùn)練加持

除了大幅降低 Sora 復(fù)現(xiàn)的技術(shù)門檻，提升視頻生成在時(shí)長(zhǎng)、分辨率、內(nèi)容等多個(gè)維度的質(zhì)量，作者團(tuán)隊(duì)還提供了 Colossal-AI 加速系統(tǒng)進(jìn)行 Sora 復(fù)現(xiàn)的高效訓(xùn)練加持。通過算子優(yōu)化和混合并行等高效訓(xùn)練策略，在處理 64 幀、512x512 分辨率視頻的訓(xùn)練中，實(shí)現(xiàn)了 1.55 倍的加速效果。同時(shí)，得益于 Colossal-AI 的異構(gòu)內(nèi)存管理系統(tǒng)，在單臺(tái)服務(wù)器上（8*H800）可以無阻礙地進(jìn)行 1 分鐘的 1080p 高清視頻訓(xùn)練任務(wù)。

此外，在作者團(tuán)隊(duì)的報(bào)告中，我們也發(fā)現(xiàn) STDiT 模型架構(gòu)在訓(xùn)練時(shí)也展現(xiàn)出卓越的高效性。和采用全注意力機(jī)制的 DiT 相比，隨著幀數(shù)的增加，STDiT 實(shí)現(xiàn)了高達(dá) 5 倍的加速效果，這在處理長(zhǎng)視頻序列等現(xiàn)實(shí)任務(wù)中尤為關(guān)鍵。

歡迎持續(xù)關(guān)注 Open-Sora 開源項(xiàng)目：https://github.com/hpcaitech/Open-Sora

作者團(tuán)隊(duì)表示，他們將會(huì)繼續(xù)維護(hù)和優(yōu)化 Open-Sora 項(xiàng)目，預(yù)計(jì)將使用更多的視頻訓(xùn)練數(shù)據(jù)，以生成更高質(zhì)量、更長(zhǎng)時(shí)長(zhǎng)的視頻內(nèi)容，并支持多分辨率特性，切實(shí)推進(jìn) AI 技術(shù)在電影、游戲、廣告等領(lǐng)域的落地。

參考鏈接：

[1] https://arxiv.org/abs/2212.09748 Scalable Diffusion Models with Transformers。

[2] https://arxiv.org/abs/2310.00426 PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis。

[3] https://arxiv.org/abs/2311.15127 Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets。

[4] https://arxiv.org/abs/2401.03048 Latte: Latent Diffusion Transformer for Video Generation。

[5] https://huggingface.co/stabilityai/sd-vae-ft-mse-original。

[6] https://github.com/google-research/text-to-text-transfer-transformer。

[7] https://github.com/haotian-liu/LLaVA。

[8] https://hpc-ai.com/blog/open-sora-v1.0。