全球首個(gè)開(kāi)源的類 Sora 架構(gòu)視頻生成模型，來(lái)了！整個(gè)訓(xùn)練流程，包括數(shù)據(jù)處理、所有訓(xùn)練細(xì)節(jié)和模型權(quán)重，全部開(kāi)放。這就是剛剛發(fā)布的 Open-Sora 1.0。
它帶來(lái)的實(shí)際效果如下，能生成繁華都市夜景中的車水馬龍。
還能用航拍視角，展現(xiàn)懸崖海岸邊，海水拍打著巖石的畫面。
亦或是延時(shí)攝影下的浩瀚星空。
自 Sora 發(fā)布以來(lái)，由于效果驚艷但技術(shù)細(xì)節(jié)寥寥，揭秘、復(fù)現(xiàn) Sora 成為了開(kāi)發(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)開(kāi)源。
那么問(wèn)題來(lái)了，復(fù)現(xiàn) Sora 如何實(shí)現(xiàn)？
Open-Sora 開(kāi)源地址：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ì)量開(kāi)源文生圖模型 PixArt-α 為基座，在此基礎(chǔ)上引入時(shí)間注意力層，將其擴(kuò)展到視頻數(shù)據(jù)上。
具體來(lái)看，整個(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í)間注意力模塊之后，交叉注意力模塊用于對(duì)齊文本的語(yǔ)意。
與全注意力機(jī)制相比，這樣的結(jié)構(gòu)大大降低了訓(xùn)練和推理開(kāi)銷。
與同樣使用空間-時(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è)階段：
大規(guī)模圖像預(yù)訓(xùn)練；
大規(guī)模視頻預(yù)訓(xùn)練；
高質(zhì)量視頻數(shù)據(jù)微調(diào)。
每個(gè)階段都會(huì)基于前一個(gè)階段的權(quán)重繼續(xù)訓(xùn)練。
相比于從零開(kāi)始單階段訓(xùn)練，多階段訓(xùn)練通過(guò)逐步擴(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í)，由于目前沒(méi)有高質(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 的開(kāi)源權(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è)量級(jí)，但是視頻的時(shí)長(zhǎng)、分辨率和質(zhì)量都更高。
通過(guò)這種方式進(jìn)行微調(diào)，能實(shí)現(xiàn)視頻生成從短到長(zhǎng)、從低分辨率到高分辨率、從低保真度到高保真度的高效擴(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)過(guò)初步估算，整個(gè)訓(xùn)練方案成功把 Open-Sora 復(fù)現(xiàn)流程控制在了 1 萬(wàn)美元左右。
數(shù)據(jù)預(yù)處理
為了進(jìn)一步降低 Sora 復(fù)現(xiàn)的門檻和復(fù)雜度，Colossal-AI 團(tuán)隊(duì)在代碼倉(cāng)庫(kù)中還提供了便捷的視頻數(shù)據(jù)預(yù)處理腳本，讓大家可以輕松啟動(dòng) Sora 復(fù)現(xiàn)預(yù)訓(xùn)練。
包括公開(kāi)視頻數(shù)據(jù)集下載、長(zhǎng)視頻根據(jù)鏡頭連續(xù)性分割為短視頻片段、使用開(kāi)源大語(yǔ)言模型 LLaVA 生成精細(xì)的提示詞。
他們提供的批量視頻標(biāo)題生成代碼可以用兩卡 3 秒標(biāo)注一個(gè)視頻，并且質(zhì)量接近于 GPT-4V。
最終得到的視頻 / 文本對(duì)可直接用于訓(xùn)練。借助他們?cè)?GitHub 上提供的開(kāi)源代碼，可以輕松地在自己的數(shù)據(jù)集上快速生成訓(xùn)練所需的視頻 / 文本對(duì)，顯著降低了啟動(dòng) Sora 復(fù)現(xiàn)項(xiàng)目的技術(shù)門檻和前期準(zhǔn)備。
高效訓(xùn)練加持
除此之外，Colossal-AI 團(tuán)隊(duì)還提供了訓(xùn)練加速方案。
通過(guò)算子優(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）可以無(wú)阻礙地進(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 倍的加速效果，這在處理長(zhǎng)視頻序列等現(xiàn)實(shí)任務(wù)中尤為關(guān)鍵。
最后，團(tuán)隊(duì)還放出了更多 Open-Sora 的生成效果。
團(tuán)隊(duì)和量子位透露，他們將長(zhǎng)期更新優(yōu)化 Open-Sora 的相關(guān)解決方案和動(dòng)態(tài)。未來(lái)將使用更多視頻訓(xùn)練數(shù)據(jù)，以生成更高質(zhì)量、更長(zhǎng)時(shí)長(zhǎng)的視頻內(nèi)容，并支持多分辨率特性。
實(shí)際應(yīng)用方面，團(tuán)隊(duì)透露將推進(jìn)在電影、游戲、廣告等領(lǐng)域落地。
感興趣的開(kāi)發(fā)者們，可訪問(wèn) GitHub 項(xiàng)目了解更多~
Open-Sora 開(kāi)源地址：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
本文來(lái)自微信公眾號(hào)：量子位 （ID：QbitAI），作者：明敏
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。