機器之心發佈
機器之心編輯部
不久前 OpenAI Sora 以其驚人的視訊生成效果迅速走紅，在一眾文生視訊模型中突出重圍，成為全球矚目的焦點。繼 2 周前推出成本直降 46% 的 Sora 訓練推理復現流程後，Colossal-AI 團隊全面開源全球首個類 Sora 架構視訊生成模型「Open-Sora 1.0」，涵蓋了整個訓練流程，包括資料處理、所有訓練細節和模型權重，攜手全球 AI 熱愛者共同推進視訊創作的新紀元。
先睹為快，我們先看一段由 Colossal-AI 團隊發佈的「Open-Sora 1.0」模型生成的都市繁華掠影視訊。
Open-Sora 1.0 生成的都市繁華掠影
這僅僅是 Sora 復現技術冰山的一角，關於以上文生視訊的模型架構、訓練好的模型權重、復現的所有訓練細節、資料預處理過程、demo 展示和詳細的上手教程，Colossal-AI 團隊已經全面免費開源在 GitHub，同時筆者第一時間聯繫了該團隊，了解到他們將不斷更新 Open-Sora 的相關解決方案和最新動態，感興趣的朋友可以持續關注 Open-Sora 的開源社區。
Open-Sora 開源地址：https://github.com/hpcaitech/Open-Sora
全面解讀 Sora 復現方案
接下來，我們將深入解讀 Sora 復現方案的多個關鍵維度，包括模型架構設計、訓練復現方案、資料預處理、模型生成效果展示以及高效訓練最佳化策略。
模型架構設計
模型採用了目前火熱的 Diffusion Transformer (DiT) [1] 架構。作者團隊以同樣使用 DiT 架構的高質量開源文生圖模型 PixArt-α [2] 為基座，在此基礎上引入時間注意力層，將其擴展到了視訊資料上。具體來說，整個架構包括一個預訓練好的 VAE，一個文字編碼器，和一個利用空間 – 時間注意力機制的 STDiT (Spatial Temporal Diffusion Transformer) 模型。其中，STDiT 每層的結構如下圖所示。它採用序列的方式在二維的空間注意力模組上疊加一維的時間注意力模組，用於建模時序關係。在時間注意力模組之後，交叉注意力模組用於對齊文字的語意。與全注意力機制相比，這樣的結構大大降低了訓練和推理開銷。與同樣使用空間 – 時間注意力機制的 Latte [3] 模型相比，STDiT 可以更好的利用已經預訓練好的圖像 DiT 的權重，從而在視訊資料上繼續訓練。
STDiT 結構示意圖
整個模型的訓練和推理流程如下。據了解，在訓練階段首先採用預訓練好的 Variational Autoencoder (VAE) 的編碼器將視訊資料進行壓縮，然後在壓縮之後的潛在空間中與文字嵌入 (text embedding) 一起訓練 STDiT 擴散模型。在推理階段，從 VAE 的潛在空間中隨機取樣出一個高斯噪聲，與提示詞嵌入 (prompt embedding) 一起輸入到 STDiT 中，得到去噪之後的特徵，最後輸入到 VAE 的解碼器，解碼得到視訊。
模型的訓練流程
訓練復現方案
我們向該團隊了解到，Open-Sora 的復現方案參考了 Stable Video Diffusion (SVD)[3] 工作，共包括三個階段，分別是：
1) 大規模圖像預訓練；
2) 大規模視訊預訓練；
3) 高質量視訊資料微調。
每個階段都會基於前一個階段的權重繼續訓練。相比於從零開始單階段訓練，多階段訓練通過逐步擴展資料，更高效地達成高質量視訊生成的目標。
訓練方案三階段
第一階段：大規模圖像預訓練
第一階段通過大規模圖像預訓練，藉助成熟的文生圖模型，有效降低視訊預訓練成本。
作者團隊向我們透露，通過網際網路上豐富的大規模圖像資料和先進的文生圖技術，我們可以訓練一個高質量的文生圖模型，該模型將作為下一階段視訊預訓練的初始化權重。同時，由於目前沒有高質量的時空 VAE，他們採用了 Stable Diffusion [5] 模型預訓練好的圖像 VAE。該策略不僅保障了初始模型的優越性能，還顯著降低了視訊預訓練的整體成本。
第二階段：大規模視訊預訓練
第二階段執行大規模視訊預訓練，增加模型泛化能力，有效掌握視訊的時間序列關聯。
我們了解到，這個階段需要使用大量視訊資料訓練，保證視訊題材的多樣性，從而增加模型的泛化能力。第二階段的模型在第一階段文生圖模型的基礎上加入了時序注意力模組，用於學習視訊中的時序關係。其餘模組與第一階段保持一致，並載入第一階段權重作為初始化，同時初始化時序注意力模組輸出為零，以達到更高效更快速的收斂。Colossal-AI 團隊使用了 PixArt-alpha [2] 的開源權重作為第二階段 STDiT 模型的初始化，以及採用了 T5 [6] 模型作為文字編碼器。同時他們採用了 256×256 的小解析度進行預訓練，進一步增加了收斂速度，降低訓練成本。
第三階段：高質量視訊資料微調
第三階段對高質量視訊資料進行微調，顯著提升視訊生成的質量。
作者團隊提及第三階段用到的視訊資料規模比第二階段要少一個量級，但是視訊的時長、解析度和質量都更高。通過這種方式進行微調，他們實現了視訊生成從短到長、從低解析度到高解析度、從低保真度到高保真度的高效擴展。
作者團隊表示，在 Open-Sora 的復現流程中，他們使用了 64 塊 H800 進行訓練。第二階段的訓練量一共是 2808 GPU hours，約合 7000 美元，第三階段的訓練量是 1920 GPU hours，大約 4500 美元。經過初步估算，整個訓練方案成功把 Open-Sora 復現流程控制在了 1 萬美元左右。
資料預處理
為了進一步降低 Sora 復現的門檻和複雜度，Colossal-AI 團隊在程式碼倉庫中還提供了便捷的視訊資料預處理腳本，讓大家可以輕鬆啟動 Sora 復現預訓練，包括公開視訊資料集下載，長視訊根據鏡頭連續性分割為短視訊片段，使用開源大語言模型 LLaVA [7] 生成精細的提示詞。作者團隊提到他們提供的批量視訊標題生成程式碼可以用兩卡 3 秒標註一個視訊，並且質量接近於 GPT-4V。最終得到的視訊 / 文字對可直接用於訓練。藉助他們在 GitHub 上提供的開源程式碼，我們可以輕鬆地在自己的資料集上快速生成訓練所需的視訊 / 文字對，顯著降低了啟動 Sora 復現項目的技術門檻和前期準備。
基於資料預處理腳本自動生成的視訊 / 文字對
模型生成效果展示
下面我們來看一下 Open-Sora 實際視訊生成效果。比如讓 Open-Sora 生成一段在懸崖海岸邊，海水拍打著岩石的航拍畫面。
再讓 Open-Sora 去捕捉山川瀑布從懸崖上澎湃而下，最終匯入湖泊的宏偉鳥瞰畫面。
除了上天還能入海，簡單輸入 prompt，讓 Open-Sora 生成了一段水中世界的鏡頭，鏡頭中一隻海龜在珊瑚礁間悠然遊弋。
Open-Sora 還能通過延時攝影的手法，向我們展現了繁星閃爍的銀河。
如果你還有更多視訊生成的有趣想法，可以訪問 Open-Sora 開源社區獲取模型權重進行免費的體驗。連結：https://github.com/hpcaitech/Open-Sora
值得注意的是，作者團隊在 Github 上提到目前版本僅使用了 400K 的訓練資料，模型的生成質量和遵循文字的能力都有待提升。例如在上面的烏龜視訊中，生成的烏龜多了一隻腳。Open-Sora 1.0 也並不擅長生成人像和複雜畫面。作者團隊在 Github 上列舉了一系列待做規劃，旨在不斷解決現有缺陷，提升生成質量。
高效訓練加持
除了大幅降低 Sora 復現的技術門檻，提升視訊生成在時長、解析度、內容等多個維度的質量，作者團隊還提供了 Colossal-AI 加速系統進行 Sora 復現的高效訓練加持。通過運算元最佳化和混合並行等高效訓練策略，在處理 64 幀、512×512 解析度視訊的訓練中，實現了1.55 倍的加速效果。同時，得益於 Colossal-AI 的異構記憶體管理系統，在單臺伺服器上（8*H800）可以無阻礙地進行 1 分鐘的 1080p 高畫質視訊訓練任務。
此外，在作者團隊的報告中，我們也發現 STDiT 模型架構在訓練時也展現出卓越的高效性。和採用全注意力機制的 DiT 相比，隨著幀數的增加，STDiT 實現了高達 5 倍的加速效果，這在處理長視訊序列等現實任務中尤為關鍵。
一覽 Open-Sora 模型視訊生成效果
最後，讓我們一睹Open-Sora模型在視訊生成上的精彩表現。
歡迎持續關注 Open-Sora 開源項目：https://github.com/hpcaitech/Open-Sora
作者團隊表示，他們將會繼續維護和最佳化 Open-Sora 項目，預計將使用更多的視訊訓練資料，以生成更高質量、更長時長的視訊內容，並支持多解析度特性，切實推進 AI 技術在電影、遊戲、廣告等領域的落地。
參考連結：
[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