全球首個(gè)開源的類Sora架構(gòu)視頻生成模型,，來了,！

整個(gè)訓(xùn)練流程，包括數(shù)據(jù)處理,、所有訓(xùn)練細(xì)節(jié)和模型權(quán)重,，全部開放。

這就是剛剛發(fā)布的Open-Sora 1.0,。

它帶來的實(shí)際效果如下,，能生成繁華都市夜景中的車水馬龍。

還能用航拍視角,，展現(xiàn)懸崖海岸邊,，海水拍打著巖石的畫面。

亦或是延時(shí)攝影下的浩瀚星空,。

自Sora發(fā)布以來,，由于效果驚艷但技術(shù)細(xì)節(jié)寥寥，揭秘,、復(fù)現(xiàn)Sora成為了開發(fā)社區(qū)最熱議話題之一,。比如Colossal-AI團(tuán)隊(duì)推出成本直降46%的Sora訓(xùn)練推理復(fù)現(xiàn)流程,。

短短兩周時(shí)間后，該團(tuán)隊(duì)再次發(fā)布最新進(jìn)展,，復(fù)現(xiàn)類Sora方案,，并將技術(shù)方案及詳細(xì)上手教程在GitHub上免費(fèi)開源,。

那么問題來了,，復(fù)現(xiàn)Sora如何實(shí)現(xiàn),？

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

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

Sora復(fù)現(xiàn)方案包括四個(gè)方面：

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

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

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

• 高效訓(xùn)練優(yōu)化策略

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

模型采用了Sora同源架構(gòu)Diffusion Transformer (DiT) ,。

它以采用DiT架構(gòu)的高質(zhì)量開源文生圖模型PixArt-α為基座，在此基礎(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í)間注意力模塊之后,，交叉注意力模塊用于對齊文本的語意。

與全注意力機(jī)制相比,，這樣的結(jié)構(gòu)大大降低了訓(xùn)練和推理開銷,。

與同樣使用空間-時(shí)間注意力機(jī)制的 Latte模型相比，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)方案

在訓(xùn)練復(fù)現(xiàn)部分，Open-Sora參考了Stable Video Diffusion (SVD),。

一共分為3個(gè)階段：

1. 大規(guī)模圖像預(yù)訓(xùn)練,；

2. 大規(guī)模視頻預(yù)訓(xùn)練,；

3. 高質(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)練。

團(tuán)隊(duì)利用互聯(lián)網(wǎng)上豐富的圖像數(shù)據(jù)和文生圖技術(shù),，先訓(xùn)練出一個(gè)高質(zhì)量的文生圖模型,，將該模型作為下一階段視頻預(yù)訓(xùn)練的初始化權(quán)重。

同時(shí),，由于目前沒有高質(zhì)量的時(shí)空VAE,，他們采用Stable Diffusion預(yù)訓(xùn)練好的圖像VAE。

這樣不僅能保障初始模型的優(yōu)越性能,，還能顯著降低視頻預(yù)訓(xùn)練的整體成本,。

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

這一階段主要增加模型的泛化能力,，有效掌握視頻的時(shí)間序列關(guān)聯(lián),。

它需要使用大量視頻數(shù)據(jù)訓(xùn)練，并且保障視頻素材的多樣性,。

同時(shí),，第二階段的模型在第一階段文生圖模型的基礎(chǔ)上加入了時(shí)序注意力模塊，用于學(xué)習(xí)視頻中的時(shí)序關(guān)系,。其余模塊與第一階段保持一致,，并加載第一階段權(quán)重作為初始化，同時(shí)初始化時(shí)序注意力模塊輸出為零,，以達(dá)到更高效更快速的收斂,。

Colossal-AI團(tuán)隊(duì)使用了PixArt-alpha的開源權(quán)重作為第二階段STDiT模型的初始化，以及采用了T5模型作為文本編碼器,。他們采用了256x256的小分辨率進(jìn)行預(yù)訓(xùn)練,，進(jìn)一步增加了收斂速度，降低訓(xùn)練成本,。

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

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

據(jù)介紹,，這一階段能顯著提升模型的生成質(zhì)量,。使用的數(shù)據(jù)規(guī)模比上一階段降低一個(gè)量級,，但是視頻的時(shí)長、分辨率和質(zhì)量都更高,。

通過這種方式進(jìn)行微調(diào),，能實(shí)現(xiàn)視頻生成從短到長、從低分辨率到高分辨率,、從低保真度到高保真度的高效擴(kuò)展。

值得一提的是,，Colossal-AI還詳細(xì)透露了每階段的資源使用情況,。

在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ù)集下載,、長視頻根據(jù)鏡頭連續(xù)性分割為短視頻片段、使用開源大語言模型LLaVA生成精細(xì)的提示詞,。

他們提供的批量視頻標(biāo)題生成代碼可以用兩卡 3 秒標(biāo)注一個(gè)視頻,，并且質(zhì)量接近于 GPT-4V。

最終得到的視頻/文本對可直接用于訓(xùn)練,。借助他們在GitHub上提供的開源代碼,，可以輕松地在自己的數(shù)據(jù)集上快速生成訓(xùn)練所需的視頻/文本對，顯著降低了啟動(dòng)Sora復(fù)現(xiàn)項(xiàng)目的技術(shù)門檻和前期準(zhǔn)備,。

高效訓(xùn)練加持

除此之外,，Colossal-AI團(tuán)隊(duì)還提供了訓(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ù)器上（8H800）可以無阻礙地進(jìn)行1分鐘的1080p高清視頻訓(xùn)練任務(wù),。

而且團(tuán)隊(duì)還發(fā)現(xiàn)STDiT模型架構(gòu)在訓(xùn)練時(shí)也展現(xiàn)出卓越的高效性,。

和采用全注意力機(jī)制的DiT相比,，隨著幀數(shù)的增加，STDiT實(shí)現(xiàn)了高達(dá)5倍的加速效果,，這在處理長視頻序列等現(xiàn)實(shí)任務(wù)中尤為關(guān)鍵,。

最后，團(tuán)隊(duì)還放出了更多Open-Sora的生成效果,。

團(tuán)隊(duì)和量子位透露,，他們將長期更新優(yōu)化Open-Sora的相關(guān)解決方案和動(dòng)態(tài)。未來將使用更多視頻訓(xùn)練數(shù)據(jù),，以生成更高質(zhì)量,、更長時(shí)長的視頻內(nèi)容，并支持多分辨率特性,。

實(shí)際應(yīng)用方面,，團(tuán)隊(duì)透露將推進(jìn)在電影、游戲,、廣告等領(lǐng)域落地,。

感興趣的開發(fā)者們，可訪問GitHub項(xiàng)目了解更多~

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

參考鏈接：
[1]https:///abs/2212.09748 Scalable Diffusion Models with Transformers
[2]https:///abs/2310.00426 PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis
[3]https:///abs/2311.15127 Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets
[4]https:///abs/2401.03048 Latte: Latent Diffusion Transformer for Video Generation
[5]https:///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:///blog/open-sora-v1.0

點(diǎn)這里??關(guān)注我,，記得標(biāo)星哦～