原文鏈接：
硬核解讀Stable Diffusion（完整版）
2022年可謂是AIGC（AI Generated Content）元年，上半年有文生圖大模型DALL-E2和Stable Diffusion，下半年有OpenAI的文本對話大模型ChatGPT問世，這讓冷卻的AI又沸騰起來了，因為AIGC能讓更多的人真真切切感受到AI的力量。這篇文章將介紹比較火的文生圖模型Stable Diffusion（簡稱SD），Stable Diffusion不僅是一個完全開源的模型（代碼，數(shù)據(jù)，模型全部開源），而且是它的參數(shù)量只有1B左右，大部分人可以在普通的顯卡上進行推理甚至精調(diào)模型。毫不夸張的說，Stable Diffusion的出現(xiàn)和開源對AIGC的火熱和發(fā)展是有巨大推動作用的，因為它讓更多的人能快地上手AI作畫。這里將基于Hugging Face的diffusers庫深入講解SD的技術(shù)原理以及部分的實現(xiàn)細節(jié)，然后也會介紹SD的常用功能，注意本文主要以SD V1.5版本為例，在最后也會簡單介紹 SD 2.0版本以及基于SD的擴展應用。
SD模型原理
SD是CompVis、Stability AI和LAION等公司研發(fā)的一個文生圖模型，它的模型和代碼是開源的，而且訓練數(shù)據(jù)LAION-5B也是開源的。SD在開源90天github倉庫就收獲了33K的stars，可見這個模型是多受歡迎。
SD是一個基于latent的擴散模型，它在UNet中引入text condition來實現(xiàn)基于文本生成圖像。SD的核心來源于Latent Diffusion這個工作，常規(guī)的擴散模型是基于pixel的生成模型，而Latent Diffusion是基于latent的生成模型，它先采用一個autoencoder將圖像壓縮到latent空間，然后用擴散模型來生成圖像的latents，最后送入autoencoder的decoder模塊就可以得到生成的圖像。基于latent的擴散模型的優(yōu)勢在于計算效率更高效，因為圖像的latent空間要比圖像pixel空間要小，這也是SD的核心優(yōu)勢。文生圖模型往往參數(shù)量比較大，基于pixel的方法往往限于算力只生成64x64大小的圖像，比如OpenAI的DALL-E2和谷歌的Imagen，然后再通過超分辨模型將圖像分辨率提升至256x256和1024x1024；而基于latent的SD是在latent空間操作的，它可以直接生成256x256和512x512甚至更高分辨率的圖像。
SD模型的主體結(jié)構(gòu)如下圖所示，主要包括三個模型：
-
autoencoder：encoder將圖像壓縮到latent空間，而decoder將latent解碼為圖像；
-
CLIP text encoder：提取輸入text的text embeddings，通過cross attention方式送入擴散模型的UNet中作為condition；
-
UNet：擴散模型的主體，用來實現(xiàn)文本引導下的latent生成。
對于SD模型，其autoencoder模型參數(shù)大小為84M，CLIP text encoder模型大小為123M，而UNet參數(shù)大小為860M，所以SD模型的總參數(shù)量約為1B。
autoencoder
autoencoder是一個基于encoder-decoder架構(gòu)的圖像壓縮模型，對于一個大小為的輸入圖像，encoder模塊將其編碼為一個大小為的latent，其中為下采樣率（downsampling factor）。在訓練autoencoder過程中，除了采用L1重建損失外，還增加了感知損失（perceptual loss，即LPIPS，具體見論文The Unreasonable Effectiveness of Deep Features as a Perceptual Metric）以及基于patch的對抗訓練。輔助loss主要是為了確保重建的圖像局部真實性以及避免模糊，具體損失函數(shù)見latent diffusion的loss部分。同時為了防止得到的latent的標準差過大，采用了兩種正則化方法：第一種是KL-reg，類似VAE增加一個latent和標準正態(tài)分布的KL loss，不過這里為了保證重建效果，采用比較小的權(quán)重（～10e-6）；第二種是VQ-reg，引入一個VQ （vector quantization）layer，此時的模型可以看成是一個VQ-GAN，不過VQ層是在decoder模塊中，這里VQ的codebook采樣較高的維度（8192）來降低正則化對重建效果的影響。latent diffusion論文中實驗了不同參數(shù)下的autoencoder模型，如下表所示，可以看到當較小和較大時，重建效果越好（PSNR越大），這也比較符合預期，畢竟此時壓縮率小。
論文進一步將不同的autoencoder在擴散模型上進行實驗，在ImageNet數(shù)據(jù)集上訓練同樣的步數(shù)（2M steps），其訓練過程的生成質(zhì)量如下所示，可以看到過小的（比如1和2）下收斂速度慢，此時圖像的感知壓縮率較小，擴散模型需要較長的學習；而過大的其生成質(zhì)量較差，此時壓縮損失過大。
當在4～16時，可以取得相對好的效果。SD采用基于KL-reg的autoencoder，其中下采樣率，特征維度為，當輸入圖像為512x512大小時將得到64x64x4大小的latent。autoencoder模型時在OpenImages數(shù)據(jù)集上基于256x256大小訓練的，但是由于autoencoder的模型是全卷積結(jié)構(gòu)的（基于ResnetBlock），所以它可以擴展應用在尺寸>256的圖像上。下面我們給出使用diffusers庫來加載autoencoder模型，并使用autoencoder來實現(xiàn)圖像的壓縮和重建，代碼如下所示：
import?torch
from?diffusers?import?AutoencoderKL
import?numpy?as?np
from?PIL?import?Image
#加載模型:?autoencoder可以通過SD權(quán)重指定subfolder來單獨加載
autoencoder?=?AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="vae")
autoencoder.to("cuda",?dtype=torch.float16)
#?讀取圖像并預處理
raw_image?=?Image.open("boy.png").convert("RGB").resize((256,?256))
image?=?np.array(raw_image).astype(np.float32)?/?127.5?-?1.0
image?=?image[None].transpose(0,?3,?1,?2)
image?=?torch.from_numpy(image)
#?壓縮圖像為latent并重建
with?torch.inference_mode():
????latent?=?autoencoder.encode(image.to("cuda",?dtype=torch.float16)).latent_dist.sample()
????rec_image?=?autoencoder.decode(latent).sample
????rec_image?=?(rec_image?/?2?+?0.5).clamp(0,?1)
????rec_image?=?rec_image.cpu().permute(0,?2,?3,?1).numpy()
????rec_image?=?(rec_image?*?255).round().astype("uint8")
????rec_image?=?Image.fromarray(rec_image[0])
rec_image
這里我們給出了兩張圖片在256x256和512x512下的重建效果對比，如下所示，第一列為原始圖片，第二列為512x512尺寸下的重建圖，第三列為256x256尺寸下的重建圖。對比可以看出，autoencoder將圖片壓縮到latent后再重建其實是有損的，比如會出現(xiàn)文字和人臉的畸變，在256x256分辨率下是比較明顯的，512x512下效果會好很多。
這種有損壓縮肯定是對SD的生成圖像質(zhì)量是有一定影響的，不過好在SD模型基本上是在512x512以上分辨率下使用的。為了改善這種畸變，stabilityai在發(fā)布SD 2.0時同時發(fā)布了兩個在LAION子數(shù)據(jù)集上精調(diào)的autoencoder，注意這里只精調(diào)autoencoder的decoder部分，SD的UNet在訓練過程只需要encoder部分，所以這樣精調(diào)后的autoencoder可以直接用在先前訓練好的UNet上（這種技巧還是比較通用的，比如谷歌的Parti也是在訓練好后自回歸生成模型后，擴大并精調(diào)ViT-VQGAN的decoder模塊來提升生成質(zhì)量）。我們也可以直接在diffusers中使用這些autoencoder，比如mse版本（采用mse損失來finetune的模型）：
autoencoder?=?AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse/")
對于同樣的兩張圖，這個mse版本的重建效果如下所示，可以看到相比原始版本的autoencoder，畸變是有一定改善的。
由于SD采用的autoencoder是基于KL-reg的，所以這個autoencoder在編碼圖像時其實得到的是一個高斯分布DiagonalGaussianDistribution（分布的均值和標準差），然后通過調(diào)用sample方法來采樣一個具體的latent（調(diào)用mode方法可以得到均值）。由于KL-reg的權(quán)重系數(shù)非常小，實際得到latent的標準差還是比較大的，latent diffusion論文中提出了一種rescaling方法：首先計算出第一個batch數(shù)據(jù)中的latent的標準差，然后采用的系數(shù)來rescale latent，這樣就盡量保證latent的標準差接近1（防止擴散過程的SNR較高，影響生成效果，具體見latent diffusion論文的D1部分討論），然后擴散模型也是應用在rescaling的latent上，在解碼時只需要將生成的latent除以，然后再送入autoencoder的decoder即可。對于SD所使用的autoencoder，這個rescaling系數(shù)為0.18215。
CLIP text encoder
SD采用CLIP text encoder來對輸入text提取text embeddings，具體的是采用目前OpenAI所開源的最大CLIP模型：clip-vit-large-patch14，這個CLIP的text encoder是一個transformer模型（只有encoder模塊）：層數(shù)為12，特征維度為768，模型參數(shù)大小是123M。對于輸入text，送入CLIP text encoder后得到最后的hidden states（即最后一個transformer block得到的特征），其特征維度大小為77x768（77是token的數(shù)量），這個細粒度的text embeddings將以cross attention的方式送入UNet中。在transofmers庫中，可以如下使用CLIP text encoder：
from?transformers?import?CLIPTextModel,?CLIPTokenizer
text_encoder?=?CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="text_encoder").to("cuda")
#?text_encoder?=?CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")
tokenizer?=?CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="tokenizer")
#?tokenizer?=?CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
#?對輸入的text進行tokenize，得到對應的token?ids
prompt?=?"a?photograph?of?an?astronaut?riding?a?horse"
text_input_ids?=?text_tokenizer(
????prompt,
????padding="max_length",
????max_length=tokenizer.model_max_length,
????truncation=True,
????return_tensors="pt"
).input_ids
#?將token?ids送入text?model得到77x768的特征
text_embeddings?=?text_encoder(text_input_ids.to("cuda"))[0]
值得注意的是，這里的tokenizer最大長度為77（CLIP訓練時所采用的設置），當輸入text的tokens數(shù)量超過77后，將進行截斷，如果不足則進行paddings，這樣將保證無論輸入任何長度的文本（甚至是空文本）都得到77x768大小的特征。在訓練SD的過程中，CLIP text encoder模型是凍結(jié)的。在早期的工作中，比如OpenAI的GLIDE和latent diffusion中的LDM均采用一個隨機初始化的tranformer模型來提取text的特征，但是最新的工作都是采用預訓練好的text model。比如谷歌的Imagen采用純文本模型T5 encoder來提出文本特征，而SD則采用CLIP text encoder，預訓練好的模型往往已經(jīng)在大規(guī)模數(shù)據(jù)集上進行了訓練，它們要比直接采用一個從零訓練好的模型要好。
UNet
SD的擴散模型是一個860M的UNet，其主要結(jié)構(gòu)如下圖所示（這里以輸入的latent為64x64x4維度為例），其中encoder部分包括3個CrossAttnDownBlock2D模塊和1個DownBlock2D模塊，而decoder部分包括1個UpBlock2D模塊和3個CrossAttnUpBlock2D模塊，中間還有一個UNetMidBlock2DCrossAttn模塊。encoder和decoder兩個部分是完全對應的，中間存在skip connection。注意3個CrossAttnDownBlock2D模塊最后均有一個2x的downsample操作，而DownBlock2D模塊是不包含下采樣的。
文章來源：http://www.zghlxwxcb.cn/news/detail-446631.html
其中CrossAttnDownBlock2D模塊的主要結(jié)構(gòu)如下圖所示，text condition將通過CrossAttention模塊嵌入進來，此時Attention的query是UNet的中間特征，而key和value則是text embeddings。SD和DDPM一樣采用預測noise的方法來訓練UNet，其訓練損失也和DDPM一樣：這里的為text embeddings，此時的模型是一個條件擴散模型?；赿iffusers庫，我們可以很快實現(xiàn)SD的訓練，其核心代碼如下所示（這里參考diffusers庫下examples中的finetune代碼）：文章來源地址http://www.zghlxwxcb.cn/news/detail-446631.html
import?torch
from?diffusers?import?AutoencoderKL,?UNet2DConditionModel,?DDPMScheduler
from?transformers?import?CLIPTextModel,?CLIPTokenizer
import?torch.nn.functional?as?F
#?加載autoencoder
vae?=?AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="vae")
#?加載text?encoder
text_encoder?=?CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="text_encoder")
tokenizer?=?CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5",?subfolder="tokenizer")
#?初始化UNet
unet?=?UNet2DConditionModel(**model_config)?#?model_config為模型參數(shù)配置
#?定義scheduler
noise_scheduler?=?DDPMScheduler(
????beta_start=0.00085,?beta_end=0.012,?beta_schedule="scaled_linear",?num_train_timesteps=1000
)
#?凍結(jié)vae和text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
opt?=?torch.optim.AdamW(unet.parameters(),?lr=1e-4)
for?step,?batch?in?enumerate(train_dataloader):
????with?torch.no_grad():
????????#?將image轉(zhuǎn)到latent空間
????????latents?=?vae.encode(batch["image"]).latent_dist.sample()
????????latents?=?latents?*?vae.config.scaling_factor?#?rescaling?latents
????????#?提取text?embeddings
????????text_input_ids?=?text_tokenizer(
????????????batch["text"],
????????????padding="max_length",
????????????max_length=tokenizer.model_max_length,
????????????truncation=True,
????????????return_tensors="pt"
??).input_ids
??text_embeddings?=?text_encoder(text_input_ids)[0]
????
????#?隨機采樣噪音
????noise?=?torch.ra到了這里，關(guān)于強大到離譜！硬核解讀Stable Diffusion（完整版）的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
