明敏 發自 凹非寺
量子位 | 公眾號 qbitai
全球首個開源的類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模型作為文本編碼器。他們採用了256x256的小解析度進行預訓練，進一步增加了收斂速度，降低訓練成本。
△open-sora生成效果（提示詞：水中世界的鏡頭，鏡頭中一隻海龜在珊瑚礁間悠然游弋）
第三階段是高質量視頻數據微調。
據介紹，這一階段能顯著提升模型的生成質量。使用的數據規模比上一階段降低一個量級，但是視頻的時長、解析度和質量都更高。
通過這種方式進行微調，能實現視頻生成從短到長、從低解析度到高解析度、從低保真度到高保真度的高效擴展。
值得一提的是，colossal-ai還詳細透露了每階段的資源使用情況。
在open-sora的復現流程中，他們使用了64塊h800進行訓練。第二階段的訓練量一共是 2808 gpu hours，約合7000美元，第三階段的訓練量是1920 gpu hours，大約4500美元。經過初步估算，整個訓練方案成功把open-sora復現流程式控制制在了1萬美元左右。
數據預處理
為了進一步降低sora復現的門檻和複雜度，colossal-ai團隊在代碼倉庫中還提供了便捷的視頻數據預處理腳本，讓大家可以輕鬆啟動sora復現預訓練。
包括公開視頻數據集下載、長視頻根據鏡頭連續性分割為短視頻片段、使用開源大語言模型llava生成精細的提示詞。
他們提供的批量視頻標題生成代碼可以用兩卡 3 秒標註一個視頻，並且質量接近於 gpt-4v。
最終得到的視頻/文本對可直接用於訓練。藉助他們在github上提供的開源代碼，可以輕鬆地在自己的數據集上快速生成訓練所需的視頻/文本對，顯著降低了啟動sora復現項目的技術門檻和前期準備。
高效訓練加持
除此之外，colossal-ai團隊還提供了訓練加速方案。
通過運算元優化和混合併行等高效訓練策略，在處理64幀、512x512解析度視頻的訓練中，實現了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
