全球首個開源的類Sora架構視訊生成模型，來了！整個訓練流程，包括資料處理、所有訓練細節和模型權重，全部開放。這就是剛剛發布的Open-Sora 1.0。它帶來的實際效果如下，能生成繁華都市夜景中的車水馬龍。

還能用空拍視角，展現懸崖海岸邊，海水拍打著岩石的畫面。

亦或是縮時攝影下的浩瀚星空。

自Sora發布以來，由於效果驚艷但技術細節寥寥，揭秘、復現Sora成為了開發社群最熱門話題之一。例如Colossal-AI團隊推出成本直降46%的Sora訓練推理復現流程。

短短兩週後，團隊再次發布最新進展，復現類Sora方案，並將技術方案及詳細上手教學在GitHub上免費開源。

那麼問題來了，復現Sora如何實現？

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

全面解讀Sora復現方案

Sora復現方案包括四個面向：

模型架構設計

訓練復現方案

資料預處理

高效訓練優化策略

模型架構設計

模型採用了Sora同源架構Diffusion Transformer (DiT) 。

它以採用DiT架構的高品質開源文生圖模型PixArt-α為基座，在此基礎上引入時間注意力層，將其擴展到視訊資料上。

具體來看，整個架構包括一個預先訓練好的VAE，一個文字編碼器和一個利用空間-時間注意力機制的STDiT (Spatial Temporal Diffusion Transformer)模型。

其中，STDiT 每層的結構如下圖所示。

它採用串列的方式在二維的空間注意力模組上疊加一維的時間注意力模組，用於建模時序關係。在時間注意力模組之後，交叉注意力模組用於對齊文字的語義。

與全注意力機制相比，這樣的結構大大降低了訓練和推理開銷。

與同樣使用空間-時間注意力機制的Latte模型相比，STDiT 可以更好的利用已經預先訓練好的影像DiT 的權重，從而在視訊資料上繼續訓練。

STDiT結構示意圖

整個模型的訓練和推理流程如下。

據了解，在訓練階段首先採用預先訓練好的Variational Autoencoder (VAE)的編碼器將視訊資料進行壓縮，然後在壓縮之後的潛在空間中與文字嵌入(text embedding)一起訓練STDiT擴散模型。

在推理階段，從VAE的潛在空間中隨機採樣出一個高斯噪聲，與提示詞嵌入(prompt embedding)一起輸入到STDiT中，得到去噪之後的特徵，最後輸入到VAE的解碼器，解碼得到視頻。

模型訓練流程訓練復現方案

在訓練復現部分，Open-Sora參考了Stable Video Diffusion (SVD)。

一共分為3個階段：

大規模影像預訓練；

大規模視訊預訓練；

高品質視訊數據微調。

每個階段都會基於前一個階段的權重繼續訓練。

相較於從零開始單階段訓練，多階段訓練透過逐步擴展數據，更有效率地達成高品質視訊生成的目標。

訓練方案三階段

第一階段是大規模影像預訓練。

團隊利用網路上豐富的影像資料和文生圖技術，先訓練出一個高品質的文生圖模型，將模型作為下一階段影片預訓練的初始化權重。

同時，由於目前沒有高品質的時空VAE，他們採用Stable Diffusion預先訓練好的影像VAE。

這樣不僅能保障初始模型的優越性能，還能顯著降低視訊預訓練的整體成本。

第二階段是大規模視訊預訓練。

此階段主要增加模型的泛化能力，有效掌握影片的時間序列關聯。

它需要使用大量視訊資料訓練，並且保障視訊素材的多樣性。

同時，第二階段的模型在第一階段文生圖模型的基礎上加入了時序注意力模組，用於學習影片中的時序關係。其餘模組與第一階段保持一致，並載入第一階段權重作為初始化，同時初始化時序注意力模組輸出為零，以達到更有效率、更快速的收斂。

Colossal-AI團隊使用了PixArt-alpha的開源權重作為第二階段STDiT模型的初始化，以及採用了T5模型作為文字編碼器。他們採用了256×256的小解析度進行預訓練，進一步增加了收斂速度，降低訓練成本。

Open-Sora生成效果（提示詞：水中世界的鏡頭，鏡頭中一隻海龜在珊瑚礁間悠然游弋）

第三階段是高品質視訊資料微調。

據介紹，此階段能顯著提升模型的生成品質。使用的資料規模比上一階段降低一個量級，但是影片的時長、解析度和品質都更高。

透過這種方式進行微調，能實現視訊生成從短到長、從低解析度到高解析度、從低保真度到高保真度的高效擴展。

值得一提的是，Colossal-AI也詳細透露了每階段的資源運用。

在Open-Sora的復現流程中，他們使用了64塊H800進行訓練。第二階段的訓練量總共是2808 GPU hours，約7,000美元，第三階段的訓練量是1,920 GPU hours，約4,500美元。經過初步估算，整個訓練計畫成功把Open-Sora復現流程控制在了1萬美元左右。

資料預處理

為了進一步降低Sora復現的門檻和複雜度，Colossal-AI團隊在程式碼倉庫中也提供了便捷的視訊資料預處理腳本，讓大家可以輕鬆啟動Sora復現預訓練。

包括公開視訊資料集下載、長影片依據鏡頭連續性分割為短影片片段、使用開源大語言模型LLaVA產生精細的提示字。

他們提供的批量視頻標題生成代碼可以用兩卡3 秒標註一個視頻，並且質量接近於GPT-4V。

最終得到的視訊/文字對可直接用於訓練。借助他們在GitHub上提供的開源程式碼，可以輕鬆地在自己的資料集上快速產生訓練所需的影片/文字對，顯著降低了啟動Sora復現專案的技術門檻和前期準備。

高效訓練加持

除此之外，Colossal-AI團隊也提供了訓練加速方案。

透過算子優化和混合平行等高效訓練策略，在處理64幀、512×512解析度視訊的訓練中，實現了1.55倍的加速效果。

同時，由於Colossal-AI的異質記憶體管理系統，在單一伺服器上（8H800）可以無阻礙地進行1分鐘的1080p高清視訊訓練任務。

而且團隊也發現STDiT模型架構在訓練時也展現出卓越的高效能。

和採用全注意力機制的DiT相比，隨著幀數的增加，STDiT實現了高達5倍的加速效果，這在處理長視訊序列等現實任務中尤其關鍵。

最後，團隊還放出了更多Open-Sora的生成效果。

團隊和量子位元透露，他們將長期更新優化Open-Sora的相關解決方案和動態。未來將使用更多影片訓練數據，以產生更高品質、更長時長的影片內容，並支援多重解析度特性。

實際應用方面，團隊透露將推動在電影、遊戲、廣告等領域落地。

有興趣的開發者們，可造訪GitHub專案了解更多~

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

參考連結：

[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

來源：量子位