別等OpenAI了,全球首個類Sora搶先開源!所有訓練細節/模型權重全公開,成本僅1萬美元
不久前OpenAI Sora以其驚人的視訊生成效果迅速走紅,在一眾文生視訊模型中突出重圍,成為全球矚目的焦點。
繼2週前推出成本直降46%的Sora訓練推理復現流程後,Colossal-AI團隊全面開源全球首個類Sora架構視頻生成模型“Open-Sora 1.0” ——涵蓋了整個訓練流程,包括數據處理、所有訓練細節與模型權重,攜手全球AI熱愛者共同推動影片創作的新紀元。
Open-Sora開源位址:https://github.com/hpcaitech/Open-Sora
先睹為快,我們先看一段由Colossal-AI團隊發布的「Open-Sora 1.0」模型生成的都市繁華掠影影片。
Open-Sora 1.0生成的都市繁華掠影
這只是Sora復現技術冰山的一角,關於以上文生影片的模型架構、訓練好的模型權重、復現的所有訓練細節、資料預處理過程、demo展示和詳細的上手教程,Colossal-AI團隊已經全面免費開源在GitHub。
同時,新智元第一時間聯繫了團隊,了解到他們將不斷更新Open-Sora的相關解決方案和最新動態,有興趣的朋友可以持續關注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]模型作為文本編碼器。同時他們採用了256x256的小解析度進行預訓練,進一步增加了收斂速度,降低訓練成本。
第三階段:高品質視訊資料微調
第三阶段对高质量视频数据进行微调,显著提升视频生成的质量。
作者團隊提及第三階段用到的影片資料規模比第二階段少一個量級,但是影片的長度、解析度和品質都更高。透過這種方式進行微調,他們實現了視訊生成從短到長、從低解析度到高解析度、從低保真度到高保真度的高效擴展。
作者團隊表示,在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幀、512x512解析度視訊的訓練中,實現了1.55倍的加速效果。
同時,由於Colossal-AI的異質記憶體管理系統,在單一伺服器上(8 x H800)可以無阻礙地進行1分鐘的1080p高清視訊訓練任務。
此外,在作者團隊的報告中,我們也發現STDiT模型架構在訓練時也展現出卓越的高效性。
和採用全注意力機制的DiT相比,隨著幀數的增加,STDiT實現了高達5倍的加速效果,這在處理長視訊序列等現實任務中尤其關鍵。
一覽Open-Sora模型影片生成效果
歡迎持續關注Open-Sora開源專案:https://github.com/hpcaitech/Open-Sora
作者團隊提及,他們將會繼續維護和優化Open-Sora項目,預計將使用更多的視頻訓練數據,以產生更高質量、更長時長的視頻內容,並支持多分辨率特性,切實推進AI技術在電影、遊戲、廣告等領域的落地。