SD技術(shù)原理以及部分實現(xiàn)細節(jié)

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的擴展應(yīng)用。

SD模型原理

SD是CompVis、Stability AI和LAION等公司研發(fā)的一個文生圖模型，它的模型和代碼是開源的，而且訓(xùn)練數(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)文本引導(dǎo)下的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）。在訓(xùn)練autoencoder過程中，除了采用L1重建損失 外，還增加了感知損失 （perceptual loss，即LPIPS，具體見論文The Unreasonable Effectiveness of Deep Features as a Perceptual Metric）以及基于patch的對抗訓(xùn)練 。輔助loss主要是為了確保重建的圖像局部真實性以及避免模糊，具體損失函數(shù)見latent diffusion的loss部分。同時為了防止得到的latent的標(biāo)準(zhǔn)差過大，采用了兩種正則化方法：第一種是KL-reg ，類似VAE增加一個latent和標(biāo)準(zhǔn)正態(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模型，如下表所示，可以看到當(dāng)較小和較大時，重建效果越好（PSNR越大），這也比較符合預(yù)期，畢竟此時壓縮率小。

論文進一步將不同的autoencoder在擴散模型上進行實驗，在ImageNet數(shù)據(jù)集上訓(xùn)練同樣的步數(shù)（2M steps），其訓(xùn)練過程的生成質(zhì)量如下所示，可以看到過小的（比如1和2）下收斂速度慢，此時圖像的感知壓縮率較小，擴散模型需要較長的學(xué)習(xí)；而過大的其生成質(zhì)量較差，此時壓縮損失過大。

當(dāng)在4～16時，可以取得相對好的效果。SD采用基于KL-reg的autoencoder，其中下采樣率，特征維度為，當(dāng)輸入圖像為512x512大小時將得到64x64x4大小的latent。autoencoder模型時在OpenImages數(shù)據(jù)集上基于256x256大小訓(xùn)練的，但是由于autoencoder的模型是全卷積結(jié)構(gòu)的（基于ResnetBlock），所以它可以擴展應(yīng)用在尺寸>256的圖像上。下面我們給出使用diffusers庫來加載autoencoder模型，并使用autoencoder來實現(xiàn)圖像的壓縮和重建，代碼如下所示：

importtorch
fromdiffusersimportAutoencoderKL
importnumpyasnp
fromPILimportImage

#加載模型:autoencoder可以通過SD權(quán)重指定subfolder來單獨加載
autoencoder=AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder="vae")
autoencoder.to("cuda",dtype=torch.float16)


>基于SpringBoot+MyBatisPlus+Vue&Element實現(xiàn)的后臺管理系統(tǒng)+用戶小程序，支持RBAC動態(tài)權(quán)限、多租戶、數(shù)據(jù)權(quán)限、工作流、三方登錄、支付、短信、商城等功能
>
>*項目地址：
>*視頻教程：

#讀取圖像并預(yù)處理
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)


>基于SpringCloudAlibaba+Gateway+Nacos+RocketMQ+Vue&Element實現(xiàn)的后臺管理系統(tǒng)+用戶小程序，支持RBAC動態(tài)權(quán)限、多租戶、數(shù)據(jù)權(quán)限、工作流、三方登錄、支付、短信、商城等功能
>
>*項目地址：
>*視頻教程：

#壓縮圖像為latent并重建
withtorch.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在訓(xùn)練過程只需要encoder部分，所以這樣精調(diào)后的autoencoder可以直接用在先前訓(xùn)練好的UNet上（這種技巧還是比較通用的，比如谷歌的Parti也是在訓(xùn)練好后自回歸生成模型后，擴大并精調(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（分布的均值和標(biāo)準(zhǔn)差），然后通過調(diào)用sample方法來采樣一個具體的latent（調(diào)用mode方法可以得到均值）。由于KL-reg的權(quán)重系數(shù)非常小，實際得到latent的標(biāo)準(zhǔn)差還是比較大的，latent diffusion論文中提出了一種rescaling方法：首先計算出第一個batch數(shù)據(jù)中的latent的標(biāo)準(zhǔn)差，然后采用的系數(shù)來rescale latent，這樣就盡量保證latent的標(biāo)準(zhǔn)差接近1（防止擴散過程的SNR較高，影響生成效果，具體見latent diffusion論文的D1部分討論），然后擴散模型也是應(yīng)用在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：

fromtransformersimportCLIPTextModel,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，得到對應(yīng)的tokenids
prompt="aphotographofanastronautridingahorse"
text_input_ids=text_tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).input_ids

#將tokenids送入textmodel得到77x768的特征
text_embeddings=text_encoder(text_input_ids.to("cuda"))[0]

值得注意的是，這里的tokenizer最大長度為77（CLIP訓(xùn)練時所采用的設(shè)置），當(dāng)輸入text的tokens數(shù)量超過77后，將進行截斷，如果不足則進行paddings，這樣將保證無論輸入任何長度的文本（甚至是空文本）都得到77x768大小的特征。在訓(xùn)練SD的過程中，CLIP text encoder模型是凍結(jié)的 。在早期的工作中，比如OpenAI的GLIDE和latent diffusion中的LDM均采用一個隨機初始化的tranformer模型來提取text的特征，但是最新的工作都是采用預(yù)訓(xùn)練好的text model。比如谷歌的Imagen采用純文本模型T5 encoder來提出文本特征，而SD則采用CLIP text encoder，預(yù)訓(xùn)練好的模型往往已經(jīng)在大規(guī)模數(shù)據(jù)集上進行了訓(xùn)練，它們要比直接采用一個從零訓(xùn)練好的模型要好。

UNet

SD的擴散模型是一個860M的UNet，其主要結(jié)構(gòu)如下圖所示（這里以輸入的latent為64x64x4維度為例），其中encoder部分包括3個CrossAttnDownBlock2D模塊和1個DownBlock2D模塊，而decoder部分包括1個UpBlock2D模塊和3個CrossAttnUpBlock2D模塊，中間還有一個UNetMidBlock2DCrossAttn模塊。encoder和decoder兩個部分是完全對應(yīng)的，中間存在skip connection。注意3個CrossAttnDownBlock2D模塊最后均有一個2x的downsample操作，而DownBlock2D模塊是不包含下采樣的。

其中CrossAttnDownBlock2D模塊的主要結(jié)構(gòu)如下圖所示，text condition將通過CrossAttention模塊嵌入進來，此時Attention的query是UNet的中間特征，而key和value則是text embeddings。

SD和DDPM一樣采用預(yù)測noise的方法來訓(xùn)練UNet，其訓(xùn)練損失也和DDPM一樣：這里的為text embeddings，此時的模型是一個條件擴散模型?；赿iffusers庫，我們可以很快實現(xiàn)SD的訓(xùn)練，其核心代碼如下所示（這里參考diffusers庫下examples中的finetune代碼）：

importtorch
fromdiffusersimportAutoencoderKL,UNet2DConditionModel,DDPMScheduler
fromtransformersimportCLIPTextModel,CLIPTokenizer
importtorch.nn.functionalasF

#加載autoencoder
vae=AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder="vae")
#加載textencoder
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)

forstep,batchinenumerate(train_dataloader):
withtorch.no_grad():
#將image轉(zhuǎn)到latent空間
latents=vae.encode(batch["image"]).latent_dist.sample()
latents=latents*vae.config.scaling_factor#rescalinglatents
#提取textembeddings
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.randn_like(latents)
bsz=latents.shape[0]
#隨機采樣timestep
timesteps=torch.randint(0,noise_scheduler.num_train_timesteps,(bsz,),device=latents.device)
timesteps=timesteps.long()

#將noise添加到latent上，即擴散過程
noisy_latents=noise_scheduler.add_noise(latents,noise,timesteps)

#預(yù)測noise并計算loss
model_pred=unet(noisy_latents,timesteps,encoder_hidden_states=text_embeddings).sample
loss=F.mse_loss(model_pred.float(),noise.float(),reduction="mean")

opt.step()
opt.zero_grad()

注意的是SD的noise scheduler雖然也是采用一個1000步長的scheduler，但是不是linear的，而是scaled linear，具體的計算如下所示：

betas=torch.linspace(beta_start**0.5,beta_end**0.5,num_train_timesteps,dtype=torch.float32)**2

在訓(xùn)練條件擴散模型時，往往會采用Classifier-Free Guidance （這里簡稱為CFG），所謂的CFG簡單來說就是在訓(xùn)練條件擴散模型的同時也訓(xùn)練一個無條件的擴散模型，同時在采樣階段將條件控制下預(yù)測的噪音和無條件下的預(yù)測噪音組合在一起來確定最終的噪音，具體的計算公式如下所示：

這里的為guidance scale ，當(dāng)越大時，condition起的作用越大，即生成的圖像其更和輸入文本一致。CFG的具體實現(xiàn)非常簡單，在訓(xùn)練過程中，我們只需要以一定的概率（比如10%）隨機drop掉text 即可，這里我們可以將text置為空字符串（前面說過此時依然能夠提取text embeddings）。這里并沒有介紹CLF背后的技術(shù)原理，感興趣的可以閱讀CFG的論文Classifier-Free Diffusion Guidance以及guided diffusion的論文Diffusion Models Beat GANs on Image Synthesis。CFG對于提升條件擴散模型的圖像生成效果是至關(guān)重要的 。

訓(xùn)練細節(jié)

前面我們介紹了SD的模型結(jié)構(gòu)，這里我們也簡單介紹一下SD的訓(xùn)練細節(jié)，主要包括訓(xùn)練數(shù)據(jù)和訓(xùn)練資源，這方面也是在SD的Model Card上有說明。首先是訓(xùn)練數(shù)據(jù)，SD在laion2B-en 數(shù)據(jù)集 上訓(xùn)練的，它是laion-5b 數(shù)據(jù)集 的一個子集，更具體的說它是laion-5b中的英文（文本為英文）數(shù)據(jù)集。laion-5b數(shù)據(jù)集是從網(wǎng)頁數(shù)據(jù)Common Crawl中篩選出來的圖像-文本對數(shù)據(jù)集，它包含5.85B的圖像-文本對，其中文本為英文的數(shù)據(jù)量為2.32B，這就是laion2B-en數(shù)據(jù)集。

下面是laion2B-en數(shù)據(jù)集的元信息（圖片width和height，以及文本長度）統(tǒng)計分析：其中圖片的width和height均在256以上的樣本量為1324M，在512以上的樣本量為488M，而在1024以上的樣本為76M；文本的平均長度為67。

laion數(shù)據(jù)集中除了圖片（下載URL，圖像width和height）和文本（描述文本）的元信息外，還包含以下信息：

similarity：使用CLIP ViT-B/32計算出來的圖像和文本余弦相似度；

pwatermark：使用一個圖片水印檢測器檢測的概率值，表示圖片含有水印的概率；

punsafe：圖片是否安全，或者圖片是不是NSFW，使用基于CLIP的檢測器來估計；

AESTHETIC_SCORE：圖片的美學(xué)評分（1-10），這個是后來追加的，首先選擇一小部分圖片數(shù)據(jù)集讓人對圖片的美學(xué)打分，然后基于這個標(biāo)注數(shù)據(jù)集來訓(xùn)練一個打分模型，并對所有樣本計算估計的美學(xué)評分。

上面是laion數(shù)據(jù)集的情況，下面我們來介紹SD訓(xùn)練數(shù)據(jù)集的具體情況，SD的訓(xùn)練是多階段的 （先在256x256尺寸上預(yù)訓(xùn)練，然后在512x512尺寸上精調(diào)），不同的階段產(chǎn)生了不同的版本：

SD v1.1：在laion2B-en數(shù)據(jù)集上以256x256大小訓(xùn)練237,000步，上面我們已經(jīng)說了，laion2B-en數(shù)據(jù)集中256以上的樣本量共1324M；然后在laion5B的高分辨率數(shù)據(jù)集以512x512尺寸訓(xùn)練194,000步，這里的高分辨率數(shù)據(jù)集是圖像尺寸在1024x1024以上，共170M樣本。

SD v1.2：以SD v1.1為初始權(quán)重，在improved_aesthetics_5plus數(shù)據(jù)集上以512x512尺寸訓(xùn)練515,000步數(shù)，這個improved_aesthetics_5plus數(shù)據(jù)集上laion2B-en數(shù)據(jù)集中美學(xué)評分在5分以上的子集（共約600M樣本），注意這里過濾了含有水印的圖片（pwatermark>0.5)以及圖片尺寸在512x512以下的樣本。

SD v1.3：以SD v1.2為初始權(quán)重，在improved_aesthetics_5plus數(shù)據(jù)集上繼續(xù)以512x512尺寸訓(xùn)練195,000步數(shù)，不過這里采用了CFG（以10%的概率隨機drop掉text）。

SD v1.4：以SD v1.2為初始權(quán)重，在improved_aesthetics_5plus數(shù)據(jù)集上采用CFG以512x512尺寸訓(xùn)練225,000步數(shù)。

SD v1.5：以SD v1.2為初始權(quán)重，在improved_aesthetics_5plus數(shù)據(jù)集上采用CFG以512x512尺寸訓(xùn)練595,000步數(shù)。

其實可以看到SD v1.3、SD v1.4和SD v1.5其實是以SD v1.2為起點在improved_aesthetics_5plus數(shù)據(jù)集上采用CFG訓(xùn)練過程中的不同checkpoints，目前最常用的版本是SD v1.4和SD v1.5 。SD的訓(xùn)練是采用了32臺8卡的A100機器 （32 x 8 x A100_40GB GPUs），所需要的訓(xùn)練硬件還是比較多的，但是相比語言大模型還好。單卡的訓(xùn)練batch size為2，并采用gradient accumulation，其中g(shù)radient accumulation steps=2，那么訓(xùn)練的總batch size就是32x8x2x2=2048 。訓(xùn)練優(yōu)化器采用AdamW ，訓(xùn)練采用warmup，在初始10,000步后學(xué)習(xí)速率升到0.0001 ，后面保持不變。至于訓(xùn)練時間，文檔上只說了用了150,000小時，這個應(yīng)該是A100卡時，如果按照256卡A100來算的話，那么大約需要訓(xùn)練25天左右 。

模型評測

上面介紹了模型訓(xùn)練細節(jié)，那么最后的問題就是模型評測了。對于文生圖模型，目前常采用的定量指標(biāo)是FID （Fréchet inception distance）和CLIP score，其中FID可以衡量生成圖像的逼真度（image fidelity），而CLIP score評測的是生成的圖像與輸入文本的一致性，其中FID越低越好，而CLIP score是越大越好。當(dāng)CFG的gudiance scale參數(shù)設(shè)置不同時，F(xiàn)ID和CLIP score會發(fā)生變化，下圖為不同的gudiance scale參數(shù)下，SD模型在COCO2017驗證集上的評測結(jié)果，注意這里是zero-shot評測，即SD模型并沒有在COCO訓(xùn)練數(shù)據(jù)集上精調(diào)。

可以看到當(dāng)gudiance scale=3時，F(xiàn)ID最低；而當(dāng)gudiance scale越大時，CLIP score越大，但是FID同時也變大。在實際應(yīng)用時，往往會采用較大的gudiance scale，比如SD模型默認采用7.5，此時生成的圖像和文本有較好的一致性。從不同版本的對比曲線上看，SD的采用CFG訓(xùn)練后三個版本其實差別并沒有那么大，其中SD v1.5相對好一點，但是明顯要未采用CFG訓(xùn)練的版本要好的多，這說明CFG訓(xùn)練是比較關(guān)鍵的。目前在模型對比上，大家往往是比較不同模型在COCO驗證集上的zero-shot FID-30K（選擇30K的樣本），大家往往就選擇模型所能得到的最小FID來比較，下面為eDiff和GigaGAN兩篇論文所報道的不同文生圖模型的FID對比（由于SD并沒有給出FID-30K，所以大家應(yīng)該都是自己用開源SD的模型計算的，由于選擇樣本不同，可能結(jié)果存在差異）：

可以看到SD雖然FID不是最好的，但是也能達到比較低的FID（大約在8～9之間）。不過雖然學(xué)術(shù)界常采用FID來定量比較模型，但是FID有很大的局限性，它并不能很好地衡量生成圖像的質(zhì)量，也是因為這個原因，谷歌的Imagen引入了人工評價，先建立一個評測數(shù)據(jù)集DrawBench（包含200個不同類型的text），然后用不同的模型來生成圖像，讓人去評價同一個text下不同模型生成的圖像，這種評測方式比較直接，但是可能也受一些主觀因素的影響?？偠灾?，目前的評價方式都有一定的局限性，最好還是直接上手使用來比較不同的模型。

SD的主要應(yīng)用

下面來介紹SD的主要應(yīng)用，這包括文生圖 ，圖生圖 以及圖像inpainting 。其中文生圖是SD的基礎(chǔ)功能：根據(jù)輸入文本生成相應(yīng)的圖像，而圖生圖和圖像inpainting是在文生圖的基礎(chǔ)上延伸出來的兩個功能。

文生圖

根據(jù)文本生成圖像這是文生圖的最核心的功能，下圖為SD的文生圖的推理流程圖：首先根據(jù)輸入text用text encoder提取text embeddings，同時初始化一個隨機噪音noise（latent上的，512x512圖像對應(yīng)的noise維度為64x64x4），然后將text embeddings和noise送入擴散模型UNet中生成去噪后的latent，最后送入autoencoder的decoder模塊得到生成的圖像。

使用diffusers庫，我們可以直接調(diào)用StableDiffusionPipeline來實現(xiàn)文生圖，具體代碼如下所示：

importtorch
fromdiffusersimportStableDiffusionPipeline
fromPILimportImage

#組合圖像，生成grid
defimage_grid(imgs,rows,cols):
assertlen(imgs)==rows*cols

w,h=imgs[0].size
grid=Image.new('RGB',size=(cols*w,rows*h))
grid_w,grid_h=grid.size

fori,imginenumerate(imgs):
grid.paste(img,box=(i%cols*w,i//cols*h))
returngrid

#加載文生圖pipeline
pipe=StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",#或者使用SDv1.4:"CompVis/stable-diffusion-v1-4"
torch_dtype=torch.float16
).to("cuda")

#輸入text，這里text又稱為prompt
prompts=[
"aphotographofanastronautridingahorse",
"Acuteotterinarainbowwhirlpoolholdingshells,watercolor",
"Anavocadoarmchair",
"Awhitedogwearingsunglasses"
]

generator=torch.Generator("cuda").manual_seed(42)#定義隨機seed，保證可重復(fù)性

#執(zhí)行推理
images=pipe(
prompts,
height=512,
width=512,
num_inference_steps=50,
guidance_scale=7.5,
negative_prompt=None,
num_images_per_prompt=1,
generator=generator
).images

grid=image_grid(images,rows=1,cols=4)
grid

生成的圖像效果如下所示：

這里可以通過指定width和height來決定生成圖像的大小，前面說過SD最后是在512x512尺度上訓(xùn)練的，所以生成512x512尺寸效果是最好的，但是實際上SD可以生成任意尺寸的圖片：一方面autoencoder支持任意尺寸的圖片的編碼和解碼，另外一方面擴散模型UNet也是支持任意尺寸的latents生成的（UNet是卷積+attention的混合結(jié)構(gòu)）。然而，生成512x512以外的圖片會存在一些問題，比如生成低分辨率圖像時，圖像的質(zhì)量大幅度下降，下圖為同樣的文本在256x256尺寸下的生成效果：

如果是生成512x512以上分辨率的圖像，圖像質(zhì)量雖然沒問題，但是可能會出現(xiàn)重復(fù)物體以及物體被拉長的情況，下圖為分別為768x512和512x768尺寸下的生成效果，可以看到部分圖像存在一定的問題：

所以雖然SD的架構(gòu)上支持任意尺寸的圖像生成，但訓(xùn)練是在固定尺寸上（512x512），生成其它尺寸圖像還是會存在一定的問題。解決這個問題的辦法就相對比較簡單，就是采用多尺度策略訓(xùn)練，比如NovelAI提出采用Aspect Ratio Bucketing策略來在二次元數(shù)據(jù)集上精調(diào)模型，這樣得到的模型就很大程度上避免SD的這個問題，目前大部分開源的基于SD的精調(diào)模型往往都采用類似的多尺度策略來精調(diào)。比如我們采用開源的dreamlike-diffusion-1.0模型（基于SD v1.5精調(diào)的），其生成的圖像效果在變尺寸上就好很多：

另外一個參數(shù)是num_inference_steps，它是指推理過程中的去噪步數(shù)或者采樣步數(shù) 。SD在訓(xùn)練過程采用的是步數(shù)為1000的noise scheduler，但是在推理時往往采用速度更快的scheduler：只需要少量的采樣步數(shù)就能生成不錯的圖像，比如SD默認采用PNDM scheduler ，它只需要采樣50步就可以出圖。當(dāng)然我們也可以換用其它類型的scheduler，比如DDIM scheduler和DPM-Solver scheduler。我們可以在diffusers中直接替換scheduler，比如我們想使用DDIM：

fromdiffusersimportDDIMScheduler

#注意這里的clip_sample要關(guān)閉，否則生成圖像存在問題，因為不能對latent進行clip
pipe.scheduler=DDIMScheduler.from_config(pipe.scheduler.config,clip_sample=False)

換成DDIM后，同樣的采樣步數(shù)生成的圖像如下所示，在部分細節(jié)上和PNDM有差異：

當(dāng)然采樣步數(shù)越大，生成的圖像質(zhì)量越好，但是相應(yīng)的推理時間也更久 。這里我們可以試驗一下不同采樣步數(shù)下的生成效果，以宇航員騎馬為例，下圖展示了采樣步數(shù)為10，20，30，50，70和100時的生成圖像，可以看到采樣步數(shù)增加后，圖像生成質(zhì)量是有一定的提升的，當(dāng)采樣步數(shù)為30時就能生成相對穩(wěn)定的圖像。

我們要討論的第三個參數(shù)是guidance_scale，前面說過當(dāng)CFG的guidance_scale越大時，生成的圖像應(yīng)該會和輸入文本更一致，這里我們同樣以宇航員騎馬為例來測試不同guidance_scale下的圖像生成效果。下圖為guidance_scale為1，3，5，7，9和11下生成的圖像對比，可以看到當(dāng)guidance_scale較低時生成的圖像效果是比較差的，當(dāng)guidance_scale在7～9時，生成的圖像效果是可以的 ，當(dāng)采用更大的guidance_scale比如11，圖像的色彩過飽和而看起來不自然，所以SD默認采用的guidance_scale為7.5 。

過大的guidance_scale之所以出現(xiàn)問題，主要是由于訓(xùn)練和測試的不一致，過大的guidance_scale會導(dǎo)致生成的樣本超出范圍。谷歌的Imagen論文提出一種dynamic thresholding策略來解決這個問題，所謂的dynamic thresholding是相對于原來的static thresholding，static thresholding策略是直接將生成的樣本clip到[-1, 1]范圍內(nèi)（Imagen是基于pixel的擴散模型，這里是將圖像像素值歸一化到-1到1之間），但是會在過大的guidance_scale時產(chǎn)生很多的飽含像素點。而dynamic thresholding策略是先計算樣本在某個百分位下（比如99%）的像素絕對值，然后如果它超過1時就采用來進行clip，這樣就可以大大減少過飽和的像素。兩種策略的具體實現(xiàn)代碼如下所示：

dynamic thresholding策略對于Imagen是比較關(guān)鍵的，它使得Imagen可以采用較大的guidance_scale來生成更自然的圖像。下圖為兩種thresholding策略下生成圖像的對比：

雖然SD是基于latent的擴散模型，但依然可以采用類似的dynamic thresholding策略，感興趣的可以參考目前的一個開源實現(xiàn)：sd-dynamic-thresholding，使用dynamic thresholding策略后，SD可以在較大的guidance_scale下生成相對自然的圖像。

另外一個比較容易忽略的參數(shù)是negative_prompt，這個參數(shù)和CFG有關(guān)，前面說過，SD采用了CFG來提升生成圖像的質(zhì)量。使用CFG，去噪過程的噪音預(yù)測不僅僅依賴條件擴散模型，也依賴無條件擴散模型：這里的negative_prompt便是無條件擴散模型的text輸入，前面說過訓(xùn)練過程中我們將text置為空字符串來實現(xiàn)無條件擴散模型，所以這里：negative_prompt = None = ""。但是有時候我們可以使用不為空的negative_prompt來避免模型生成的圖像包含不想要的東西 ，因為從上述公式可以看到這里的無條件擴散模型是我們想遠離的部分。下面我們來舉幾個具體的例子，首先來看生成人物圖像的一個例子，這里的輸入文本為"a portrait of a beautiful blonde woman"，其生成的圖像如下所示：

可以看到生成的圖像效果并不好，比如出現(xiàn)一些臉部的畸變，但是我們可以設(shè)置negative_prompt來提升生成效果，這里我們將negative_prompt設(shè)置為"cropped, lowres, poorly drawn face, out of frame, poorly drawn hands, blurry"，這些描述都是負面的。改變negative_prompt后，生成的圖像效果有一個明顯的提升：

第二個例子是一個建筑物，這里的輸入文本為"A Hyperrealistic photograph of German architectural modern home"，默認圖像生成效果如下所示：

雖然生成的圖像效果不錯，但是如果只想要一個干凈的建筑物，而不想背景中含有樹木和草地等，此時我們可以通過設(shè)置negative prompt來達到這種效果。這里將negative prompt設(shè)為"trees, bushes, leaves, greenery"，其生成的建筑物就干凈了很多：

可以看到合理使用negative prompt能夠幫助我們?nèi)コ幌胍臇|西來提升圖像生成效果。一般情況下，輸入的text或者prompt我們稱之為“正向提示詞 ”，而negative prompt稱之為“反向提示詞 ”，想要生成的好的圖像，不僅要選擇好的正向提示詞，也需要好的反向提示詞，這和文本生成模型也比較類似：都需要好的prompt。這里也舉一個對正向prompt優(yōu)化的例子（這個例子來源于微軟的工作Optimizing Prompts for Text-to-Image Generation），這里的原始prompt為"A rabbit is wearing a space suit"，可以看到直接生成的效果其實是不盡人意的：

但是如果我們將prompt改為"A rabbit is wearing a space suit, digital Art, Greg rutkowski, Trending cinematographic artstation"，其生成的效果就大大提升：

這里我們其實只是在原有的prompt基礎(chǔ)加上了一些描述詞，有時候我們稱之為“魔咒 ”，不同的模型可能會有不同的魔咒。上述我們討論了SD的文生圖的主要參數(shù)，這里簡單總結(jié)一下：

SD默認生成512x512大小的圖像，但實際上可以生成其它分辨率的圖像，但是可能會出現(xiàn)不協(xié)調(diào)，如果采用多尺度策略訓(xùn)練，會改善這種情況；

采用快速的noise scheduler，SD在去噪步數(shù)為30～50步時就能生成穩(wěn)定的圖像；

SD的guidance_scale設(shè)置為7～9是比較穩(wěn)定的，過小和過大都會出現(xiàn)圖像質(zhì)量下降，實際使用中可以根據(jù)具體情況靈活調(diào)節(jié)；

可以使用negative prompt來去除不想要的東西來改善圖像生成效果；

好的prompt對圖像生成效果是至關(guān)重要的。

上邊我們介紹了如何使用SD進行文生圖以及一些主要參數(shù)，在最后我們也給出文生圖這個pipeline的內(nèi)部流程代碼，如下所示：

importtorch
fromdiffusersimportAutoencoderKL,UNet2DConditionModel,DDIMScheduler
fromtransformersimportCLIPTextModel,CLIPTokenizer
fromtqdm.autoimporttqdm


model_id="runwayml/stable-diffusion-v1-5"
#1.加載autoencoder
vae=AutoencoderKL.from_pretrained(model_id,subfolder="vae")
#2.加載tokenizer和textencoder
tokenizer=CLIPTokenizer.from_pretrained(model_id,subfolder="tokenizer")
text_encoder=CLIPTextModel.from_pretrained(model_id,subfolder="text_encoder")
#3.加載擴散模型UNet
unet=UNet2DConditionModel.from_pretrained(model_id,subfolder="unet")
#4.定義noisescheduler
noise_scheduler=DDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,#don'tclipsample,thex0instablediffusionnotinrange[-1,1]
set_alpha_to_one=False,
)

#將模型復(fù)制到GPU上
device="cuda"
vae.to(device,dtype=torch.float16)
text_encoder.to(device,dtype=torch.float16)
unet=unet.to(device,dtype=torch.float16)

#定義參數(shù)
prompt=[
"Adragonfruitwearingkaratebeltinthesnow",
"AsmallcactuswearingastrawhatandneonsunglassesintheSaharadesert",
"Aphotoofaraccoonwearinganastronauthelmet,lookingoutofthewindowatnight",
"Acuteotterinarainbowwhirlpoolholdingshells,watercolor"
]
height=512
width=512
num_inference_steps=50
guidance_scale=7.5
negative_prompt=""
batch_size=len(prompt)
#隨機種子
generator=torch.Generator(device).manual_seed(2023)


withtorch.no_grad():
#獲取text_embeddings
text_input=tokenizer(prompt,padding="max_length",max_length=tokenizer.model_max_length,truncation=True,return_tensors="pt")
text_embeddings=text_encoder(text_input.input_ids.to(device))[0]
#獲取unconditionaltextembeddings
max_length=text_input.input_ids.shape[-1]
uncond_input=tokenizer(
[negative_prompt]*batch_size,padding="max_length",max_length=max_length,return_tensors="pt"
)
uncond_embeddings=text_encoder(uncond_input.input_ids.to(device))[0]
#拼接為batch，方便并行計算
text_embeddings=torch.cat([uncond_embeddings,text_embeddings])

#生成latents的初始噪音
latents=torch.randn(
(batch_size,unet.in_channels,height//8,width//8),
generator=generator,device=device
)
latents=latents.to(device,dtype=torch.float16)

#設(shè)置采樣步數(shù)
noise_scheduler.set_timesteps(num_inference_steps,device=device)

#scaletheinitialnoisebythestandarddeviationrequiredbythescheduler
latents=latents*noise_scheduler.init_noise_sigma#forDDIM,init_noise_sigma=1.0

timesteps_tensor=noise_scheduler.timesteps

#Dodenoisesteps
fortintqdm(timesteps_tensor):
#這里latens擴展2份，是為了同時計算unconditionalprediction
latent_model_input=torch.cat([latents]*2)
latent_model_input=noise_scheduler.scale_model_input(latent_model_input,t)#forDDIM,donothing

#使用UNet預(yù)測噪音
noise_pred=unet(latent_model_input,t,encoder_hidden_states=text_embeddings).sample

#執(zhí)行CFG
noise_pred_uncond,noise_pred_text=noise_pred.chunk(2)
noise_pred=noise_pred_uncond+guidance_scale*(noise_pred_text-noise_pred_uncond)

#計算上一步的noisylatents：x_t->x_t-1
latents=noise_scheduler.step(noise_pred,t,latents).prev_sample

#注意要對latents進行scale
latents=1/0.18215*latents
#使用vae解碼得到圖像
image=vae.decode(latents).sample

圖生圖

圖生圖（image2image）是對文生圖功能的一個擴展 ，這個功能來源于SDEdit這個工作，其核心思路也非常簡單：給定一個筆畫的色塊圖像，可以先給它加一定的高斯噪音（執(zhí)行擴散過程）得到噪音圖像，然后基于擴散模型對這個噪音圖像進行去噪，就可以生成新的圖像，但是這個圖像在結(jié)構(gòu)和布局和輸入圖像基本一致。

對于SD來說，圖生圖的流程圖如下所示，相比文生圖流程來說，這里的初始latent不再是一個隨機噪音，而是由初始圖像經(jīng)過autoencoder編碼之后的latent加高斯噪音得到，這里的加噪過程就是擴散過程。要注意的是，去噪過程的步數(shù)要和加噪過程的步數(shù)一致，就是說你加了多少噪音，就應(yīng)該去掉多少噪音，這樣才能生成想要的無噪音圖像。

在diffusers中，我們可以使用StableDiffusionImg2ImgPipeline來實現(xiàn)文生圖，具體代碼如下所示：

importtorch
fromdiffusersimportStableDiffusionImg2ImgPipeline
fromPILimportImage

#加載圖生圖pipeline
model_id="runwayml/stable-diffusion-v1-5"
pipe=StableDiffusionImg2ImgPipeline.from_pretrained(model_id,torch_dtype=torch.float16).to("cuda")

#讀取初始圖片
init_image=Image.open("init_image.png").convert("RGB")

#推理
prompt="Afantasylandscape,trendingonartstation"
generator=torch.Generator(device="cuda").manual_seed(2023)

image=pipe(
prompt=prompt,
image=init_image,
strength=0.8,
guidance_scale=7.5,
generator=generator
).images[0]
image

相比文生圖的pipeline，圖生圖的pipeline還多了一個參數(shù)strength，這個參數(shù)介于0-1之間，表示對輸入圖片加噪音的程度，這個值越大加的噪音越多，對原始圖片的破壞也就越大，當(dāng)strength=1時，其實就變成了一個隨機噪音，此時就相當(dāng)于純粹的文生圖pipeline了。下面展示了一個具體的實例，這里的第一張圖為輸入的初始圖片，它是一個筆畫的色塊，我們可以通過圖生圖將它生成一幅具體的圖像，其中第2張圖和第3張圖的strength分別是0.5和0.8，可以看到當(dāng)strength=0.5時，生成的圖像和原圖比較一致，但是就比較簡單了，當(dāng)strength=0.8時，生成的圖像偏離原圖更多，但是圖像的質(zhì)感有一個明顯的提升。

圖生圖這個功能一個更廣泛的應(yīng)用是在風(fēng)格轉(zhuǎn)換上，比如給定一張人像，想生成動漫風(fēng)格的圖像。這里我們可以使用動漫風(fēng)格的開源模型anything-v4.0，它是基于SD v1.5在動漫風(fēng)格數(shù)據(jù)集上finetune的，使用它可以更好地利用圖生圖將人物動漫化。下面的第1張為輸入人物圖像，采用的prompt為"masterpiece, best quality, 1girl, red hair, medium hair, green eyes"，后面的圖像是strength分別為0.3-0.9下生成的圖像?？梢钥吹皆诓煌膕trength下圖像有不同的生成效果，其中strength=0.6時我覺得效果是最好的。

總結(jié)來看，圖生圖其實核心也是依賴了文生圖的能力，其中strength這個參數(shù)需要靈活調(diào)節(jié)來得到滿意的圖像 。在最后，我們也給出圖生圖pipeline的內(nèi)部主要代碼，如下所示：

importPIL
importnumpyasnp
importtorch
fromdiffusersimportAutoencoderKL,UNet2DConditionModel,DDIMScheduler
fromtransformersimportCLIPTextModel,CLIPTokenizer
fromtqdm.autoimporttqdm


model_id="runwayml/stable-diffusion-v1-5"
#1.加載autoencoder
vae=AutoencoderKL.from_pretrained(model_id,subfolder="vae")
#2.加載tokenizer和textencoder
tokenizer=CLIPTokenizer.from_pretrained(model_id,subfolder="tokenizer")
text_encoder=CLIPTextModel.from_pretrained(model_id,subfolder="text_encoder")
#3.加載擴散模型UNet
unet=UNet2DConditionModel.from_pretrained(model_id,subfolder="unet")
#4.定義noisescheduler
noise_scheduler=DDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,#don'tclipsample,thex0instablediffusionnotinrange[-1,1]
set_alpha_to_one=False,
)

#將模型復(fù)制到GPU上
device="cuda"
vae.to(device,dtype=torch.float16)
text_encoder.to(device,dtype=torch.float16)
unet=unet.to(device,dtype=torch.float16)

#預(yù)處理init_image
defpreprocess(image):
w,h=image.size
w,h=map(lambdax:x-x%32,(w,h))#resizetointegermultipleof32
image=image.resize((w,h),resample=PIL.Image.LANCZOS)
image=np.array(image).astype(np.float32)/255.0
image=image[None].transpose(0,3,1,2)
image=torch.from_numpy(image)
return2.0*image-1.0

#參數(shù)設(shè)置
prompt=["Afantasylandscape,trendingonartstation"]
num_inference_steps=50
guidance_scale=7.5
strength=0.8
batch_size=1
negative_prompt=""
generator=torch.Generator(device).manual_seed(2023)

init_image=PIL.Image.open("init_image.png").convert("RGB")

withtorch.no_grad():
#獲取prompt的text_embeddings
text_input=tokenizer(prompt,padding="max_length",max_length=tokenizer.model_max_length,truncation=True,return_tensors="pt")
text_embeddings=text_encoder(text_input.input_ids.to(device))[0]
#獲取unconditionaltextembeddings
max_length=text_input.input_ids.shape[-1]
uncond_input=tokenizer(
[negative_prompt]*batch_size,padding="max_length",max_length=max_length,return_tensors="pt"
)
uncond_embeddings=text_encoder(uncond_input.input_ids.to(device))[0]
#拼接batch
text_embeddings=torch.cat([uncond_embeddings,text_embeddings])

#設(shè)置采樣步數(shù)
noise_scheduler.set_timesteps(num_inference_steps,device=device)
#根據(jù)strength計算timesteps
init_timestep=min(int(num_inference_steps*strength),num_inference_steps)
t_start=max(num_inference_steps-init_timestep,0)
timesteps=noise_scheduler.timesteps[t_start:]


#預(yù)處理init_image
init_input=preprocess(init_image)
init_latents=vae.encode(init_input.to(device,dtype=torch.float16)).latent_dist.sample(generator)
init_latents=0.18215*init_latents

#給init_latents加噪音
noise=torch.randn(init_latents.shape,generator=generator,device=device,dtype=init_latents.dtype)
init_latents=noise_scheduler.add_noise(init_latents,noise,timesteps[:1])
latents=init_latents#作為初始latents


#Dodenoisesteps
fortintqdm(timesteps):
#這里latens擴展2份，是為了同時計算unconditionalprediction
latent_model_input=torch.cat([latents]*2)
latent_model_input=noise_scheduler.scale_model_input(latent_model_input,t)#forDDIM,donothing

#預(yù)測噪音
noise_pred=unet(latent_model_input,t,encoder_hidden_states=text_embeddings).sample

#CFG
noise_pred_uncond,noise_pred_text=noise_pred.chunk(2)
noise_pred=noise_pred_uncond+guidance_scale*(noise_pred_text-noise_pred_uncond)

#計算上一步的noisylatents：x_t->x_t-1
latents=noise_scheduler.step(noise_pred,t,latents).prev_sample

#注意要對latents進行scale
latents=1/0.18215*latents
#解碼
image=vae.decode(latents).sample

圖像inpainting

最后我們要介紹的一項功能是圖像inpainting，它和圖生圖一樣也是文生圖功能的一個擴展。SD的圖像inpainting不是用在圖像修復(fù)上，而是主要用在圖像編輯 上：給定一個輸入圖像和想要編輯的區(qū)域mask，我們想通過文生圖來編輯mask區(qū)域的內(nèi)容。SD的圖像inpainting原理可以參考論文Blended Latent Diffusion，其主要原理圖如下所示：

它和圖生圖一樣：首先將輸入圖像通過autoencoder編碼為latent，然后加入一定的高斯噪音生成noisy latent，再進行去噪生成圖像，但是這里為了保證mask以外的區(qū)域不發(fā)生變化，在去噪過程的每一步，都將擴散模型預(yù)測的noisy latent用真實圖像同level的nosiy latent替換。在diffusers中，使用StableDiffusionInpaintPipelineLegacy可以實現(xiàn)文本引導(dǎo)下的圖像inpainting，具體代碼如下所示：

importtorch
fromdiffusersimportStableDiffusionInpaintPipelineLegacy
fromPILimportImage

#加載inpaintingpipeline
model_id="runwayml/stable-diffusion-v1-5"
pipe=StableDiffusionInpaintPipelineLegacy.from_pretrained(model_id,torch_dtype=torch.float16).to("cuda")

#讀取輸入圖像和輸入mask
input_image=Image.open("overture-creations-5sI6fQgYIuo.png").resize((512,512))
input_mask=Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512,512))

#執(zhí)行推理
prompt=["amecharobotsittingonabench","acatsittingonabench"]
generator=torch.Generator("cuda").manual_seed(0)

withtorch.autocast("cuda"):
images=pipe(
prompt=prompt,
image=input_image,
mask_image=input_mask,
num_inference_steps=50,
strength=0.75,
guidance_scale=7.5,
num_images_per_prompt=1,
generator=generator,
).images

下面是一個具體的生成效果，這里我們將輸入圖像的dog換成了mecha robot或者cat，從而實現(xiàn)了圖像編輯。

要注意的是這里的參數(shù)guidance_scale也和圖生圖一樣比較重要，要生成好的圖像，需要選擇合適的guidance_scale。如果guidance_scale=0.5時，生成的圖像由于過于受到原圖干擾而產(chǎn)生一些不協(xié)調(diào)，如下所示：

合適的prompt也比較重要，比如如果我們?nèi)サ魀rompt中的"sitting on a bench"，那么編輯的圖像效果也會出現(xiàn)不協(xié)調(diào)：

無論是上面的圖生圖還是這里的圖像inpainting，我們其實并沒有去finetune SD模型，只是擴展了它的能力，但是這兩樣功能就需要精確調(diào)整參數(shù)才能得到滿意的生成效果。這里，我們也給出StableDiffusionInpaintPipelineLegacy這個pipeline內(nèi)部的核心代碼：

importPIL
importnumpyasnp
importtorch
fromdiffusersimportAutoencoderKL,UNet2DConditionModel,DDIMScheduler
fromtransformersimportCLIPTextModel,CLIPTokenizer
fromtqdm.autoimporttqdm

defpreprocess_mask(mask):
mask=mask.convert("L")
w,h=mask.size
w,h=map(lambdax:x-x%32,(w,h))#resizetointegermultipleof32
mask=mask.resize((w//8,h//8),resample=PIL.Image.NEAREST)
mask=np.array(mask).astype(np.float32)/255.0
mask=np.tile(mask,(4,1,1))
mask=mask[None].transpose(0,1,2,3)#whatdoesthisstepdo?
mask=1-mask#repaintwhite,keepblack
mask=torch.from_numpy(mask)
returnmask

defpreprocess(image):
w,h=image.size
w,h=map(lambdax:x-x%32,(w,h))#resizetointegermultipleof32
image=image.resize((w,h),resample=PIL.Image.LANCZOS)
image=np.array(image).astype(np.float32)/255.0
image=image[None].transpose(0,3,1,2)
image=torch.from_numpy(image)
return2.0*image-1.0

model_id="runwayml/stable-diffusion-v1-5"
#1.加載autoencoder
vae=AutoencoderKL.from_pretrained(model_id,subfolder="vae")
#2.加載tokenizer和textencoder
tokenizer=CLIPTokenizer.from_pretrained(model_id,subfolder="tokenizer")
text_encoder=CLIPTextModel.from_pretrained(model_id,subfolder="text_encoder")
#3.加載擴散模型UNet
unet=UNet2DConditionModel.from_pretrained(model_id,subfolder="unet")
#4.定義noisescheduler
noise_scheduler=DDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,#don'tclipsample,thex0instablediffusionnotinrange[-1,1]
set_alpha_to_one=False,
)

#將模型復(fù)制到GPU上
device="cuda"
vae.to(device,dtype=torch.float16)
text_encoder.to(device,dtype=torch.float16)
unet=unet.to(device,dtype=torch.float16)

prompt="amecharobotsittingonabench"
strength=0.75
guidance_scale=7.5
batch_size=1
num_inference_steps=50
negative_prompt=""
generator=torch.Generator(device).manual_seed(0)

withtorch.no_grad():
#獲取prompt的text_embeddings
text_input=tokenizer(prompt,padding="max_length",max_length=tokenizer.model_max_length,truncation=True,return_tensors="pt")
text_embeddings=text_encoder(text_input.input_ids.to(device))[0]
#獲取unconditionaltextembeddings
max_length=text_input.input_ids.shape[-1]
uncond_input=tokenizer(
[negative_prompt]*batch_size,padding="max_length",max_length=max_length,return_tensors="pt"
)
uncond_embeddings=text_encoder(uncond_input.input_ids.to(device))[0]
#拼接batch
text_embeddings=torch.cat([uncond_embeddings,text_embeddings])

#設(shè)置采樣步數(shù)
noise_scheduler.set_timesteps(num_inference_steps,device=device)
#根據(jù)strength計算timesteps
init_timestep=min(int(num_inference_steps*strength),num_inference_steps)
t_start=max(num_inference_steps-init_timestep,0)
timesteps=noise_scheduler.timesteps[t_start:]


#預(yù)處理init_image
init_input=preprocess(input_image)
init_latents=vae.encode(init_input.to(device,dtype=torch.float16)).latent_dist.sample(generator)
init_latents=0.18215*init_latents
init_latents=torch.cat([init_latents]*batch_size,dim=0)
init_latents_orig=init_latents
#處理mask
mask_image=preprocess_mask(input_mask)
mask_image=mask_image.to(device=device,dtype=init_latents.dtype)
mask=torch.cat([mask_image]*batch_size)

#給init_latents加噪音
noise=torch.randn(init_latents.shape,generator=generator,device=device,dtype=init_latents.dtype)
init_latents=noise_scheduler.add_noise(init_latents,noise,timesteps[:1])
latents=init_latents#作為初始latents


#Dodenoisesteps
fortintqdm(timesteps):
#這里latens擴展2份，是為了同時計算unconditionalprediction
latent_model_input=torch.cat([latents]*2)
latent_model_input=noise_scheduler.scale_model_input(latent_model_input,t)#forDDIM,donothing

#預(yù)測噪音
noise_pred=unet(latent_model_input,t,encoder_hidden_states=text_embeddings).sample

#CFG
noise_pred_uncond,noise_pred_text=noise_pred.chunk(2)
noise_pred=noise_pred_uncond+guidance_scale*(noise_pred_text-noise_pred_uncond)

#計算上一步的noisylatents：x_t->x_t-1
latents=noise_scheduler.step(noise_pred,t,latents).prev_sample

#將unmask區(qū)域替換原始圖像的nosiylatents
init_latents_proper=noise_scheduler.add_noise(init_latents_orig,noise,torch.tensor([t]))
latents=(init_latents_proper*mask)+(latents*(1-mask))

#注意要對latents進行scale
latents=1/0.18215*latents
image=vae.decode(latents).sample

另外，runwayml在發(fā)布SD 1.5版本的同時還發(fā)布了一個inpainting模型：runwayml/stable-diffusion-inpainting，與前面所講不同的是，這是一個在SD 1.2上finetune的模型 。原來SD的UNet的輸入是64x64x4，為了實現(xiàn)inpainting，現(xiàn)在給UNet的第一個卷機層增加5個channels，分別為masked圖像的latents（經(jīng)過autoencoder編碼，64x64x4）和mask圖像（直接下采樣8x，64x64x1），增加的權(quán)重填零初始化。在diffusers中，可以使用StableDiffusionInpaintPipeline來調(diào)用這個模型，具體代碼如下：

importtorch
fromdiffusersimportStableDiffusionInpaintPipeline
fromPILimportImage
fromtqdm.autoimporttqdm
importPIL

#Loadpipeline
model_id="runwayml/stable-diffusion-inpainting/"
pipe=StableDiffusionInpaintPipeline.from_pretrained(model_id,torch_dtype=torch.float16).to("cuda")

prompt=["amecharobotsittingonabench","adogsittingonabench","abench"]

generator=torch.Generator("cuda").manual_seed(2023)

input_image=Image.open("overture-creations-5sI6fQgYIuo.png").resize((512,512))
input_mask=Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512,512))

images=pipe(
prompt=prompt,
image=input_image,
mask_image=input_mask,
num_inference_steps=50,
generator=generator,
).images

其生成的效果圖如下所示：

經(jīng)過finetune的inpainting在生成細節(jié)上可能會更好，但是有可能會喪失部分文生圖的能力，而且也比較難遷移其它finetune的SD模型。

SD 2.0

SD 2.0

Stability AI公司在2022年11月（stable-diffusion-v2-release）放出了SD 2.0版本，這里我們也簡單介紹一下相比SD 1.x版本SD 2.0的具體改進點。SD 2.0相比SD 1.x版本的主要變動在于模型結(jié)構(gòu) 和訓(xùn)練數(shù)據(jù) 兩個部分。

首先是模型結(jié)構(gòu)方面，SD 1.x版本的text encoder采用的是OpenAI的CLIP ViT-L/14模型，其模型參數(shù)量為123.65M；而SD 2.0采用了更大的text encoder：基于OpenCLIP在laion-2b數(shù)據(jù)集上訓(xùn)練的CLIP ViT-H/14模型，其參數(shù)量為354.03M，相比原來的text encoder模型大了約3倍。兩個CLIP模型的對比如下所示：

可以看到CLIP ViT-H/14模型相比原來的OpenAI的L/14模型，在imagenet1K上分類準(zhǔn)確率和mscoco多模態(tài)檢索任務(wù)上均有明顯的提升，這也意味著對應(yīng)的text encoder更強，能夠抓住更準(zhǔn)確的文本語義信息。另外是一個小細節(jié)是SD 2.0提取的是text encoder倒數(shù)第二層的特征，而SD 1.x提取的是倒數(shù)第一層的特征。由于倒數(shù)第一層的特征之后就是CLIP的對比學(xué)習(xí)任務(wù)，所以倒數(shù)第一層的特征可能部分丟失細粒度語義信息，Imagen論文（見論文D.1部分）和novelai（見novelai blog）均采用了倒數(shù)第二層特征。對于UNet模型，SD 2.0相比SD 1.x幾乎沒有改變，唯一的一個小的變動是：SD 2.0不同stage的attention模塊是固定attention head dim為64，而SD 1.0則是不同stage的attention模塊采用固定attention head數(shù)量，明顯SD 2.0的這種設(shè)定更常用，但是這個變動不會影響模型參數(shù)。然后是訓(xùn)練數(shù)據(jù)，前面說過SD 1.x版本其實最后主要采用laion-2B中美學(xué)評分為5以上的子集來訓(xùn)練，而SD 2.0版本采用評分在4.5以上的子集，相當(dāng)于擴大了訓(xùn)練數(shù)據(jù)集，具體的訓(xùn)練細節(jié)見model card。另外SD 2.0除了512x512版本的模型，還包括768x768版本的模型（https://huggingface.co/stabilityai/stable-diffusion-2），所謂的768x768模型是在512x512模型基礎(chǔ)上用圖像分辨率大于768x768的子集繼續(xù)訓(xùn)練的，不過優(yōu)化目標(biāo)不再是noise_prediction，而是采用Progressive Distillation for Fast Sampling of Diffusion Models論文中所提出的 v-objective。下圖為SD 2.0和SD 1.x版本在COCO2017驗證集上評測的對比，可以看到2.0相比1.5，CLIP score有一個明顯的提升，同時FID也有一定的提升。但是正如前面所討論的，F(xiàn)ID和CLIP score這兩個指標(biāo)均有一定的局限性，所以具體效果還是上手使用來對比。

Stability AI在發(fā)布SD 2.0的同時，還發(fā)布了另外3個模型：stable-diffusion-x4-upscaler ，stable-diffusion-2-inpainting 和stable-diffusion-2-depth 。stable-diffusion-x4-upscaler是一個基于擴散模型的4x超分模型，它也是基于latent diffusion，不過這里采用的autoencoder是基于VQ-reg的，下采樣率為。在實現(xiàn)上，它是將低分辨率圖像直接和noisy latent拼接在一起送入UNet，因為autoencoder將高分辨率圖像壓縮為原來的1/4，而低分辨率圖像也為高分辨率圖像的1/4，所以低分辨率圖像的空間維度和latent是一致的。另外，這個超分模型也采用了Cascaded Diffusion Models for High Fidelity Image Generation所提出的noise conditioning augmentation，簡單來說就是在訓(xùn)練過程中給低分辨率圖像加上高斯噪音，可以通過擴散過程來實現(xiàn)，注意這里的擴散過程的scheduler與主擴散模型的scheduler可以不一樣，同時也將對應(yīng)的noise_level（對應(yīng)擴散模型的time step）通過class labels的方式送入UNet，讓UNet知道加入噪音的程度。stable-diffusion-x4-upscaler是使用LAION中>2048x2048大小的子集（10M）訓(xùn)練的，訓(xùn)練過程中采用512x512的crops來訓(xùn)練（降低顯存消耗）。SD模型可以用來生成512x512圖像，加上這個超分模型，就可以得到2048x2048大小的圖像。

在diffusers庫中，可以如下使用這個超分模型（這里的noise level是指推理時對低分辨率圖像加入噪音的程度）：

importrequests
fromPILimportImage
fromioimportBytesIO
fromdiffusersimportStableDiffusionUpscalePipeline
importtorch

#loadmodelandscheduler
model_id="stabilityai/stable-diffusion-x4-upscaler"
pipeline=StableDiffusionUpscalePipeline.from_pretrained(model_id,torch_dtype=torch.float16)
pipeline=pipeline.to("cuda")

#let'sdownloadanimage
url="https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale/low_res_cat.png"
response=requests.get(url)
low_res_img=Image.open(BytesIO(response.content)).convert("RGB")
low_res_img=low_res_img.resize((128,128))

prompt="awhitecat"

upscaled_image=pipeline(prompt=prompt,image=low_res_img,noise_level=20).images[0]
upscaled_image.save("upsampled_cat.png")

stable-diffusion-2-inpainting是圖像inpainting模型，和前面所說的runwayml/stable-diffusion-inpainting基本一樣，不過它是在SD 2.0的512x512版本上finetune的。

stable-diffusion-2-depth 是也是在SD 2.0的512x512版本上finetune的模型，它是額外增加了圖像的深度圖作為condition，這里是直接將深度圖下采樣8x，然后和nosiy latent拼接在一起送入UNet模型中。深度圖可以作為一種結(jié)構(gòu)控制，下圖展示了加入深度圖后生成的圖像效果：

你可以調(diào)用diffusers庫中的StableDiffusionDepth2ImgPipeline來實現(xiàn)基于深度圖控制的文生圖：

importtorch
importrequests
fromPILimportImage
fromdiffusersimportStableDiffusionDepth2ImgPipeline

pipe=StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
).to("cuda")

url="http://images.cocodataset.org/val2017/000000039769.jpg"
init_image=Image.open(requests.get(url,stream=True).raw)

prompt="twotigers"
n_propmt="bad,deformed,ugly,badanotomy"
image=pipe(prompt=prompt,image=init_image,negative_prompt=n_propmt,strength=0.7).images[0]

除此之外，Stability AI公司還開源了兩個加強版的autoencoder：ft-EMA和ft-MSE（前者使用L1 loss后者使用MSE loss），前面已經(jīng)說過，它們是在LAION數(shù)據(jù)集繼續(xù)finetune decoder來增強重建效果。

SD 2.1

在SD 2.0版本發(fā)布幾周后，Stability AI又發(fā)布了SD 2.1。SD 2.0在訓(xùn)練過程中采用NSFW檢測器過濾掉了可能包含色情的圖像（punsafe=0.1），但是也同時過濾了很多人像圖片，這導(dǎo)致SD 2.0在人像生成上效果可能較差，所以SD 2.1是在SD 2.0的基礎(chǔ)上放開了限制（punsafe=0.98）繼續(xù)finetune，所以增強了人像的生成效果。

和SD 2.0一樣，SD 2.1也包含兩個版本：512x512版本和768x768版本。

SD unclip

Stability AI在2023年3月份，又放出了基于SD的另外一個模型：stable-diffusion-reimagine，它可以實現(xiàn)單個圖像的變換，即image variations，目前該模型已經(jīng)在在huggingface上開源：stable-diffusion-2-1-unclip。

這個模型是借鑒了OpenAI的DALLE2（又稱unCLIP)，unCLIP是基于CLIP的image encoder提取的image embeddings作為condition來實現(xiàn)圖像的生成。

SD unCLIP是在原來的SD模型的基礎(chǔ)上增加了CLIP的image encoder的nosiy image embeddings作為condition。具體來說，它在訓(xùn)練過程中是對提取的image embeddings施加一定的高斯噪音（也是通過擴散過程），然后將noise level對應(yīng)的time embeddings和image embeddings拼接在一起，最后再以class labels的方式送入UNet。在diffusers中，你可以調(diào)用StableUnCLIPImg2ImgPipeline來實現(xiàn)圖像的變換：

importrequests
importtorch
fromPILimportImage
fromioimportBytesIO

fromdiffusersimportStableUnCLIPImg2ImgPipeline

#StarttheStableUnCLIPImagevariationspipeline
pipe=StableUnCLIPImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1-unclip",torch_dtype=torch.float16,variation="fp16"
)
pipe=pipe.to("cuda")

#GetimagefromURL
url="https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/tarsila_do_amaral.png"
response=requests.get(url)
init_image=Image.open(BytesIO(response.content)).convert("RGB")

#Pipetomakethevariation
images=pipe(init_image).images
images[0].save("tarsila_variation.png")

其實在SD unCLIP之前，已經(jīng)有Lambda Labs開源的sd-image-variations-diffusers，它是在SD 1.4的基礎(chǔ)上finetune的模型，不過實現(xiàn)方式是直接將text embeddings替換為image embeddings，這樣也同樣可以實現(xiàn)圖像的變換。

這里SD unCLIP有兩個版本：sd21-unclip-l和sd21-unclip-h，兩者分別是采用OpenAI CLIP-L和OpenCLIP-H模型的image embeddings作為condition。如果要實現(xiàn)文生圖，還需要像DALLE2那樣訓(xùn)練一個prior模型，它可以實現(xiàn)基于文本來預(yù)測對應(yīng)的image embeddings，我們將prior模型和SD unCLIP接在一起就可以實現(xiàn)文生圖了。KakaoBrain這個公司已經(jīng)開源了一個DALLE2的復(fù)現(xiàn)版本：Karlo，它是基于OpenAI CLIP-L來實現(xiàn)的，你可以基于這個模型中prior模塊加上sd21-unclip-l來實現(xiàn)文本到圖像的生成，目前這個已經(jīng)集成了在StableUnCLIPPipeline中，或者基于stablediffusion官方倉庫來實現(xiàn)。

SD的其它特色應(yīng)用

在SD模型開源之后，社區(qū)和研究機構(gòu)也基于SD實現(xiàn)了形式多樣的特色應(yīng)用，這里我們也選擇一些比較火的應(yīng)用來介紹一下。

個性化生成

個性化生成是指的生成特定的角色或者風(fēng)格，比如給定自己幾張肖像來利用SD來生成個性化頭像。在個性化生成方面，比較重要的兩個工作是英偉達的Textual Inversion 和谷歌的DreamBooth 。Textual Inversion 這個工作的核心思路是基于用戶提供的3～5張?zhí)囟ǜ拍睿ㄎ矬w或者風(fēng)格）的圖像來學(xué)習(xí)一個特定的text embeddings，實際上只用一個word embedding就足夠了。Textual Inversion 不需要finetune UNet，而且由于text embeddings較小，存儲成本很低。目前diffusers庫已經(jīng)支持textual_inversion的訓(xùn)練。

DreamBooth 原本是谷歌提出的應(yīng)用在Imagen上的個性化生成，但是它實際上也可以擴展到SD上（更新版論文已經(jīng)增加了SD）。DreamBooth首先為特定的概念尋找一個特定的描述詞[V]，這個特定的描述詞只要是稀有的就可以，然后與Textual Inversion不同的是DreamBooth需要finetune UNet，這里為了防止過擬合，增加了一個class-specific prior preservation loss（基于SD生成同class圖像加入batch里面訓(xùn)練）來進行正則化。

由于finetune了UNet，DreamBooth往往比Textual Inversion要表現(xiàn)的要好，但是DreamBooth的存儲成本較高。目前diffusers庫已經(jīng)支持dreambooth訓(xùn)練，你也可以在sd-dreambooth-library中找到其他人上傳的模型。DreamBooth和Textual Inversion是最常用的個性化生成方法，但其實除了這兩種，還有很多其它的研究工作，比如Adobe提出的Custom Diffusion ，相比DreamBooth，它只finetune了UNet的attention模塊的KV權(quán)重矩陣，同時優(yōu)化一個新概念的token。

風(fēng)格化finetune模型

SD的另外一大應(yīng)用是采用特定風(fēng)格的數(shù)據(jù)集進行finetune，這使得模型“過擬合”在特定的風(fēng)格上 。之前比較火的novelai就是基于二次元數(shù)據(jù)在SD上finetune的模型，雖然它失去了生成其它風(fēng)格圖像的能力，但是它在二次元圖像的生成效果上比原來的SD要好很多。

目前已經(jīng)有很多風(fēng)格化的模型在huggingface上開源，這里也列出一些：

更多的模型可以直接在huggingface text-to-image模型庫上找到。此外，很多基于SD進行finetune的模型開源在civitai上，你也可以在這個網(wǎng)站上找到更多風(fēng)格的模型。值得說明的一點是，目前finetune SD模型的方法主要有兩種：一種是直接finetune了UNet，但是容易過擬合，而且存儲成本；另外一種低成本的方法是基于微軟的LoRA ，LoRA本來是用于finetune語言模型的，但是現(xiàn)在已經(jīng)可以用來finetune SD模型了，具體可以見博客Using LoRA for Efficient Stable Diffusion Fine-Tuning。

圖像編輯

圖像編輯也是SD比較火的應(yīng)用方向，這里所說的圖像編輯是指的是使用SD來實現(xiàn)對圖片的局部編輯。這里列舉兩個比較好的工作：谷歌的prompt-to-prompt和加州伯克利的instruct-pix2pix。谷歌的prompt-to-prompt 的核心是基于UNet的cross attention maps來實現(xiàn)對圖像的編輯，它的好處是不需要finetune模型，但是主要用在編輯用SD生成的圖像。

谷歌后面的工作Null-text Inversion有進一步實現(xiàn)了對真實圖片的編輯：

instruct-pix2pix 這個工作基于GPT-3和prompt-to-prompt構(gòu)建了pair的數(shù)據(jù)集，然后在SD上進行finetune，它可以輸入text instruct對圖像進行編輯：

可控生成

可控生成是SD最近比較火的應(yīng)用，這主要歸功于ControlNet，基于ControlNet可以實現(xiàn)對很多種類的可控生成，比如邊緣，人體關(guān)鍵點，草圖和深度圖等等。

其實在ControlNet之前，也有一些可控生成的工作，比如stable-diffusion-2-depth 也屬于可控生成，但是都沒有太火。我覺得ControlNet之所以火，是因為這個工作直接實現(xiàn)了各種各種的可控生成，而且訓(xùn)練的ControlNet可以遷移到其它基于SD finetune的模型上（見Transfer Control to Other SD1.X Models）：

與ControlNet同期的工作還有騰訊的T2I-Adapter以及阿里的composer-page：

stable-diffusion-webui

最后要介紹的一個比較火的應(yīng)用stable-diffusion-webui其實是用來支持SD出圖的一個web工具，它算是基于gradio框架實現(xiàn)了SD的快速部署，不僅支持SD的最基礎(chǔ)的文生圖、圖生圖以及圖像inpainting功能，還支持SD的其它拓展功能，很多基于SD的拓展應(yīng)用可以用插件的方式安裝在webui上。

后話

在OpenAI最早放出DALLE2的時候，我曾被它生成的圖像所驚艷到，但是我從來沒有想到圖像生成的AIGC會如此火爆，技術(shù)的發(fā)展太快了，這得益于互聯(lián)網(wǎng)獨有的開源精神。我想，沒有SD的開源，估計這個方向可能還會沉寂一段時間。