AI繪畫Stable Diffusion原理之Autoencoder-Latent

這篇具有很好參考價(jià)值的文章主要介紹了AI繪畫Stable Diffusion原理之Autoencoder-Latent。希望對大家有所幫助。如果存在錯誤或未考慮完全的地方，請大家不吝賜教，您也可以點(diǎn)擊"舉報(bào)違法"按鈕提交疑問。

前言

傳送門：

stable diffusion：Git｜論文

stable-diffusion-webui：Git

Google Colab Notebook部署stable-diffusion-webui：Git

kaggle Notebook部署stable-diffusion-webui：Git

AI繪畫，輸入一段文本就能生成相關(guān)的圖像，stable diffusion便是其中一個重要分支。自己對其中的原理比較感興趣，因此開啟這個系列的文章來對stable diffusion的原理進(jìn)行學(xué)習(xí)（主要是針對“文生圖”[text to image]）。

上述的stable-diffusion-webui是AUTOMATIC1111開發(fā)的一套UI操作界面，可以在自己的主機(jī)上搭建，無限生成圖像（實(shí)測2080ti完全能夠勝任），如果沒有資源，可以白嫖Google Colab或者kaggle的GPU算力，其部署教程在上面?zhèn)魉烷T。

其中stable diffusion的基礎(chǔ)模型可以hugging face下載，而C站可以下載各種風(fēng)格的模型。stable diffusion有一個很大的優(yōu)勢就是基于C站中各式各樣的模型，我們可以進(jìn)行不同風(fēng)格的AI繪畫。

而這篇文章，首先對其中的一個組件進(jìn)行學(xué)習(xí)：Autoencoder/VQGANs，可以將圖像從像素空間壓縮到低維的隱空間。

原理簡介

Stable Diffusion is a latent text-to-image diffusion model。stable diffusion本質(zhì)是一種latent diffusion models（LDMs），隱向量擴(kuò)散模型。diffusion models (DMs)將圖像的形成過程分解為去噪自動編碼器(denoising autoencoders)的一系列操作，但這些都是直接在像素空間上進(jìn)行的操作，因此對于昂貴的計(jì)算資源，特別是高像素的圖像。而LDMs則是引入隱向量空間，能夠生成超高像素的圖像。

這里，我們先整體地來了解下stable diffusion的結(jié)構(gòu)組成，后面再對每個組件進(jìn)行拆開逐一理解。整體結(jié)構(gòu)如下圖[Stable Diffusion Architecture]：

文本編碼器：人類輸入的文本即prompt，經(jīng)過CLIP模型中的Text Encoder，轉(zhuǎn)化為語義向量（Token Embeddings）；
圖像生成器（Image information Creator）：U-Net、采樣器以及Autoencoder組成。由隨機(jī)生成的純噪聲向量（即下圖中的Noisey Image）開始，通過Autoencoder編碼映射到低維的隱空間，文本語義向量作為控制條件進(jìn)行指導(dǎo)，由U-Net和采樣器不斷迭代生成新的越具有豐富語義信息的隱向量，這就是擴(kuò)散過程diffusion；
圖像解碼器（Image Decoder）- Autoencoder：迭代了一定次數(shù)之后，得到了包含豐富語義信息的隱向量（Processed Image Info Tensor），低維的隱向量經(jīng)過Autoencoder解碼到原始像素；
第2步就是LDMs和DMs的區(qū)別，LDMs是在latent space進(jìn)行擴(kuò)散，而DMs則是在pixel space，這也是性能提升的關(guān)鍵。

AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

Autoencoder

[1] 論文：Taming Transformers for High-Resolution Image Synthesis

[2] Git：taming-transformers

圖片的隱空間表征從何而來：Autoencoder，既能夠?qū)D片從像素空間壓縮到隱空間，讓擴(kuò)散過程在latent space中進(jìn)行，又可以讓圖片從隱空間重建到像素空間（即圖片重建），簡化的過程如下圖所示：

其中的encoder可以將一張圖片從RGB空間即像素空間 $x\in \mathbb{R}^{H\times W \times 3}$ ，經(jīng)過encoder編碼到隱空間表征（latent representation） $\varepsilon(x)$ ；
decoder則是將隱空間表征重建到圖片RGB $\tilde{x}=D(z)=D(\varepsilon(x))$ ；
其中， $z\in \mathbb{R}^{h \times w \times c}$ ，重要的是，控制隱空間大小的是編碼器的下采樣因子（downsampling factors）： $f=H/h=W/w，f=2^m，m \in \mathbb{N}$

AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

上述僅僅是從整體架構(gòu)層面簡單地描述了圖片的隱空間與像素空間的轉(zhuǎn)換與重建過程，但其實(shí)整個過程的細(xì)節(jié)還是比較復(fù)雜的，方法是出自VQGAN $^{[1]}$ ，其結(jié)構(gòu)如下圖所示：

論文認(rèn)為高像素的圖片合成需要模型能夠理解圖片的全局組成，使得局部和全局現(xiàn)實(shí)的生成能夠保持一致。
因此，論文使用codebook來對圖片的豐富視覺組成進(jìn)行表征，而不是像素表征，codebook即是隱空間的表現(xiàn)形式。
codebook可以大大減少的圖片組成長度（相比像素），也使得能用transformer來高效地對圖片內(nèi)部的全局交互（ global interrelations）進(jìn)行建模。

AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

Codebook

給定一張圖片 $x\in \mathbb{R}^{H\times W \times 3}$ ，需要將x表征為離散空間的codebook集合 $z_q \in \mathbb{R}^{h\times w \times n_z}$ ，其中 $\cdot w$ 可以認(rèn)為是codebook中每個code的索引，而 $n_z$ 是code的維度。學(xué)習(xí)這樣的codebook表征需要以下幾個組件：

一個離散的codebook $Z=\{z_k\}^K_{k=1} \in \mathbb{R}^{n_z}$ （可以當(dāng)成embedding來理解，參數(shù)隨機(jī)初始化，參與模型訓(xùn)練，但論文對這塊沒有清晰的描述，可以去看源碼）
CNN結(jié)構(gòu)的encoder E，可以將圖片 $x$ 編碼為 $\hat{z} \in \mathbb{R}^{h\times w \times n_z}$
CNN結(jié)構(gòu)的decoder G，能夠?qū)odebook $z_q$ 重建為圖像 $\hat{x}$
quantization操作，將 $\hat{z}$ 映射到 $z_q$

具體的 $z_q$ 編碼過程為：編碼器E將x轉(zhuǎn)化為 $\hat{z}=E(x) \in \mathbb{R}^{h\times w \times n_z}$ ，然后通過element-wise quantization $q(\cdot)$ 將每個離散的code $\hat{z}_{ij} \in \mathbb{R}^{n_z}$ 編碼到距離最近的codebook entry $z_k$ （這里產(chǎn)生的最鄰近的 $z_k$ 索引即為上圖[VQGAN]的 $s_i$ ，后續(xù)會用到） AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

這部分的損失函數(shù)如下式 AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

其中 $L_{rec}=||x-\hat{x}||^2$ 為重建loss， $sg[\cdot]$ 為stop-gradient操作。由于 $z_q$ 的quantization操作是不可微分的，因此需要用到梯度拷貝（出自straight-through gradient estimator）

Discriminator

論文：Image-to-Image Translation with Conditional Adversarial Networks

Git：https://github.com/phillipi/pix2pix

使用transformer來表征圖片的隱性圖像成分的分布，需要進(jìn)一步逼近圖片壓縮的極限和學(xué)習(xí)更富含信息的codebook，因此，論文還訓(xùn)練一個patch-based的判別器D，讓它能夠區(qū)分真實(shí)和重建的圖片： AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

真實(shí)圖像和重建圖像都會經(jīng)過一個CNN結(jié)構(gòu)的Discriminator，然后得到每個patch的預(yù)估概率，模型的訓(xùn)練目標(biāo)就是讓真實(shí)圖像的預(yù)估概率盡量都為1，而重建圖像的預(yù)估概率盡量都為0，簡而言之，就是讓Discriminator能夠識別每個patch是來自真實(shí)圖像還是重建圖像，如下圖紅框部分： AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

上述這兩部分是聯(lián)合訓(xùn)練： AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

其中， $\nabla_{G_L}[\cdot]$ 是decoder最后一層網(wǎng)絡(luò)的梯度，而 $\delta=10^{-6}$ 。

Transformers

Latent Transformers.

編碼器E和解碼器G訓(xùn)練完成之后，按照上述同樣的操作，通過E和quantization操作，可以將圖片 $x$ 表征到codebook $z_q=q(E(x)) \in \mathbb{R}^{h \times w \times n_z}$ ， $\cdot w$ 可以認(rèn)為是codebook中每個code的索引 $s_i$ ，然后將二維的索引變?yōu)橐痪S的，相當(dāng)于一個code序列 $\in \{0,...,|Z|-1\}^{h \times w}$ ： AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

到這里，我們就可以按照NLP的自回歸模型“預(yù)測下一個詞”的思路來理解：給定code索引序列（上文） $s < i$ ，利用transformer來學(xué)習(xí)下一個code索引（下文）的概率分布 $p(s_i|s<i)$ ，最大化完整表征序列的似然估計(jì) $p(s)=\prod_ip(s_i|s<i)$ ： AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

Conditioned Synthesis.

在許多圖片合成任務(wù)中，往往會加入額外的信息來控制圖片的合成過程，這個額外信息稱為 $c$ ，它可以是一個對圖片的標(biāo)簽描述或者另外的圖片。那么，學(xué)習(xí)的似然估計(jì)則變?yōu)椋?img src="https://imgs.yssmx.com/Uploads/2023/07/609425-13.png" alt="AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion" referrerpolicy="no-referrer" />

機(jī)制理解

在最后，通過源碼倉庫里的兩個實(shí)操案例notebook來理解Autoencoder這些組建的工作機(jī)制。

圖像重建.

VQGAN可以將圖片輸入編碼到低維的codebook空間（隱空間），然后再對codebook空間重建為圖片的像素空間，如下圖所示。更重要的是，這個過程的中間產(chǎn)物-隱空間，相較于像素空間，能夠以很小的特征空間來表征圖片，可以遷移到attention機(jī)制底座的模型訓(xùn)練的下流任務(wù)，比如本文的主題：Stable Diffusion。

def reconstruct_with_vqgan(x, model):
  # could also use model(x) for reconstruction but use explicit encoding and decoding here
  z, _, [_, _, indices] = model.encode(x)
  print(f"VQGAN --- {model.__class__.__name__}: latent shape: {z.shape[2:]}")
  xrec = model.decode(z)
  return xrec

AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion

草圖繪畫.

這里主要是可以幫助理解VQGAN中Transformer的作用：

草圖經(jīng)過VQGAN的編碼器得到codebook索引序列c- $s_i$ （c-僅是前綴，為了與成品圖進(jìn)行區(qū)分）；
隨機(jī)生成成品圖的codebook索引序列z- $s_i$ ；
然后草圖的索引序列c- $s_i$ 作為控制條件，即上述提到Conditioned Synthesis章節(jié)中的 $c$ ，拼接在z- $s_i$ 的前面（z- $s_i$ 每次截取一段），輸入到Transformer，去預(yù)測z- $s_i$ 的每一個位置，預(yù)測得到的索引逐步替代隨機(jī)生成的索引序列；
最后，這個生成的索引序列再進(jìn)入解碼器G重建為圖片（成品圖）。

AI繪畫Stable Diffusion原理之Autoencoder-Latent,AIGC,AIGC,stable diffusion
文章來源地址http://www.zghlxwxcb.cn/news/detail-609425.html

到了這里，關(guān)于AI繪畫Stable Diffusion原理之Autoencoder-Latent的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！