從零開始的 stable diffusion

stable diffusion 真的是橫空出世,開啓了 AIGC 的元年。不知你是否有和我一樣的困惑,這 AI 工具好像並不是那麼聽話?

前言

我們該如何才能用好 stable diffusion 這個工具呢?AI 究竟在 stable diffusion 中承擔了什麼樣的角色?如何能儘可能快、成本低地得到我們期望的結果?

源於這一系列的疑問,我開始了漫長的論文解讀。High-Resolution Image Synthesis with Latent Diffusion Models(地址:https://arxiv.org/abs/2112.10752?spm=ata.21736010.0.0.7d0b28addsl7xQ&file=2112.10752)

當然這論文看的雲裏霧裏的,加篇讀了 How does Stable Diffusion work?(地址:https://stable-diffusion-art.com/how-stable-diffusion-work/?spm=ata.21736010.0.0.7d0b28addsl7xQ)

先簡要概括下,stable diffusion 的努力基本是爲了 2 個目的:

  1. 低成本、高效驗證。設計了 Latent Space

  2. Conditioning Mechanisms。條件控制,如果不能輸出我們想要的圖片,那這就像 Monkey Coding。耗費無限的時間與資源。

這是整個內容裏最重要最核心的兩個部分。

圖片生成的幾種方式

隨着深度神經網絡的發展,生成模型已經有了巨大的發展,主流的有以下幾種:

  1. 自迴歸模型 (AutoRegressive model):按照像素點生成圖像,導致計算成本高。實驗效果還不錯

  2. 變分自編碼器 (Variational Autoencoder):Image to Latent, Latent to Image,VAE 存在生成圖像模糊或者細節問題

  3. 基於流的方法 (Glow)

  4. 生成對抗網絡 (Generative adversarial network):利用生成器(G) 與判別器 (D) 進行博弈,不斷讓生成的圖像與真實的圖像在分佈上越來越接近。

其中 AR 與 GAN 的生成都是在 pixel space 進行模型訓練與推理。

模型是如何生成圖片的?

以一隻貓作爲案例。當我們想畫一隻貓的時候,也都是從一個白板開始,框架、細節不斷完善。

對於 AI 來說,一個純是 noise 的 image 就是一個理想的白板,類似下圖展示的這樣。

從圖中的流程,我們可以看到推理的過程如下:

  1. 生成一個隨機的 noise image 圖片。這個 noise 取決於 Random 這個參數。相同的 Random 生成的 noise image 是相同的。

  2. 使用 noise predictor 預測這個圖裏加了多少 noise,生成一個 predicted noise。

  3. 使用原始的 noise 減去 predicted noise。

  4. 不斷循環 2、3,直到我們的執行 steps。

最終我們會得到一隻貓。

在這個過程中,我們會以下疑問:

  1. 如何得到一個 noise predictor?

  2. 怎麼控制我們最終能得到一隻貓?而不是一隻狗或者其他的東西?

在回答這些疑問之前,我先貼一部分公式:

我們定義一個 noise predictor:是第 t 個 step 過程中的 noise image,t 表示第 t 個 stop。

如何得到一個 noise predictor?

這是一個訓練的過程。過程如下圖所示:

  1. 選擇一張訓練用的圖片,比如說一張貓

  2. 生成一個隨機的 noise 圖片

  3. 將 noise 圖疊加到訓練用的圖片上,得到一張有一些 noise 的圖片。(這裏可以疊加 1~T 步 noise

  4. 訓練 noise predictor,告訴我們加了多少 noise。通過正確的 noise 答案來調整模型權重。

最終我們能得到一個相對準確的 noise-predictor。這是一個 U-Net model。在 stable-diffusion-model 中。

通過這一步,我們最終能得到一個 noise encoder 與 noise decoder。

PS: noise encoder 在 image2image 中會應用到。

以上 noise 與 noise-predictor 的過程均在 pixel space,那麼就會存在巨大的性能問題。比如說一張 1024x1024x3 的 RBG 圖片對應 3,145,728 個數字,需要極大的計算資源。在這裏 stable diffusion 定義了一個 Latent Space,來解決這個問題。

Latent Space

Latent Space 的提出基於一個理論:Manifold_hypothesis

它假設現實世界中,許多高維數據集實際上位於該高維空間內的低維 Latent manifolds。像 Pixel Space,就存在很多的難以感知的高頻細節,而這些都是在 Latent Space 中需要壓縮掉的信息。

那麼基於這個假設,我們先定義一個在 RGB 域的圖片

然後存在一個方法 z=varepsilon(x),,z 是 x 在 latent space 的一種表達。

這裏有一個因子 f=H/h=W/w,通常我們定義,比如說 stable-diffusion v1.5 訓練與推理圖片在 512x512x3,然後 Latent Space 的中間表達則是 4x64x64,那麼我們會有一個 decoder D 能將圖片從 Latent Space 中解碼出來。

在這個過程中我們期望,這倆圖片無限接近。

整個過程如下圖所示:

而執行這個過程的就是我們的 Variational Autoencoder,也就是 VAE。

那麼 VAE 該怎麼訓練呢?我們需要一個衡量生成圖像與訓練圖像之間的一個距離指標。

也就是

細節就不關心了,但這個指標可以用來衡量 VAE 模型的還原程度。訓練過程與 noise encoder 和 noise-predictor 非常接近。

貼一個 stable diffusion 在 FID 指標上,與其他方法的對比。下面的表格來自於無條件圖片生成。基本就是比較 Latent Space 是否有丟失重要信息。

你可能在想,爲什麼 VAE 可以把一張圖片壓縮到更小的 latent space,並且可以不丟失信息。

其實和人對圖片的理解是一樣的,自然的、優秀的圖片都不是隨機的,他們有高度的規則,比如說臉上會有眼睛、鼻子。一隻狗會有 4 條腿和一個規則的形狀。

圖像的高維性是人爲的,而自然的圖像可以很容易地壓縮爲更小的空間中而不丟失任何信息。

可能說我們修改了一張圖片的很多難以感知的細節,比如說隱藏水印,微小的亮度、對比度的修改,但修改後還是同樣的圖像嗎?我們只能說它表達的東西還是一樣的。並沒有丟失任何信息。

結合 Latent Space 與 noise predictor 的圖像生成過程

  1. 生成一個隨機的 latent space matrix,也可以叫做 latent representation。一種中間表達

  2. noise-predictor 預測這個 latent representation 的 noise. 並生成一個 latent space noise

  3. latent representation 減去 latent space noise

  4. 重複 2~3,直到 step 結束

  5. 通過 VAE 的 decoder 將 latent representation 生成最終的圖片

直到目前爲止,都還沒有條件控制的部分。按這個過程,我們最終只會得到一個隨機的圖片。

條件控制

非常關鍵,沒有條件控制,我們最終只能不斷地進行 Monkey Coding,得到源源不斷的隨機圖片。

相信你在上面的圖片生成的過程中,已經感知到一個問題了,如果只是從一堆 noise 中去掉 noise,那最後得到的爲什麼是有信息的圖片,而不是一堆 noise 呢?

noise-predictor 在訓練的時候,其實就是基於已經成像的圖片去預測 noise,那麼它預測的 noise 基本都來自於有圖像信息的訓練數據。

在這個 denoise 的過程中,noise 會被附加上各種各樣的圖像信息。

怎麼控制 noise-predictor 去選擇哪些訓練數據去預測 noise,就是條件控制的核心要素。

這裏我們以 tex2img 爲案例討論。

Text Conditioning

下面的流程圖,展示了一個 prompt 如何處理,並提供給 noise predictor。

從圖中可以看到,我們的每一個 word,都會被 tokenized。stable diffusion v1.5 使用的 openai ViT-L/14 Clip 模型來進行這個過程。

tokenized 將自然語言轉成計算機可理解的數字 (NLP),它只能將 words 轉成 token。比如說dreambeach會被 CLIP 模型拆分成dreambeach。一個 word,並不意味着一個 token。同時dreambeach也不等同於dream<space>beach,stable diffusion model 目前被限制只能使用 75 個 tokens 來進行 prompt,並不等同於 75 個 word。

同樣,這也是使用的 openai ViT-L/14 Clip model. Embedding 是一個 768 長度的向量。每一個 token 都會被轉成一個 768 長度的向量,如上案例,我們最後會得到一個4x768的矩陣。

爲什麼我們需要 embedding 呢?

比如說我們輸入了man,但這是不是同時可以意味着gentlemanguysportsmanboy。他們可能說在向量空間中,與man的距離由近而遠。而你不一定非要一個完全準確無誤的man。通過 embedding 的向量,我們可以決定究竟取多近的信息來生成圖片。對應 stable diffusion 的參數就是 (Classifier-Free Guidance scale)CFG。相當於用一個 scale 去放大距離,因此 scale 越大,對應的能獲取的信息越少,就會越遵循 prompt。而 scale 越小,則越容易獲取到關聯小,甚至無關的信息。

如何去控制 embedding?

我們經常會遇到 stable diffusion 無法準確繪製出我們想要的內容。那麼這裏我們發現了第一種條件控制的方式:textual inversion

將我們想要的 token 用一個全新的別名定義,這個別名對應一個準確的 token。那麼就能準確無誤地使用對應的 embedding 生成圖片。

這裏的 embedding 可以是新的對象,也可以是其他已存在的對象。

比如說我們用一個玩具貓訓練到 CLIP 模型中,並定義其 Tokenizer 對應的 word,同時微調 stable diffusion 的模型。而對應toy cat就能產生如下的效果。

感覺有點像 Lora 的思路,具體還得調研下 lora。

text transformer

在得到 embedding 之後,通過 text transformer 輸入給 noise-predictor

transformer 可以控制多種條件,如 class labels、image、depth map 等。

Cross-attention

具體 cross-attention 是什麼我也不是很清楚。但這裏有一個案例可以說明:

比如說我們使用 prompt "A man with blue eyes"。雖然這裏是兩個 token,但 stable diffusion 會把這兩個單詞一起成對。

這樣就能保證生成一個藍色眼睛的男人。而不是一個藍色襪子或者其他藍色信息的男人。

(cross-attention between the prompt and the image)

LoRA models modify the cross-attention module to change styles。後面在研究 Lora,這裏把原話摘到這。

感覺更像是存在blueeyes,然後有一個集合同時滿足blueeye。去取這個交叉的集合。問題:對應的 embedding 是不是不一樣的?該如何區分blue planet in eyeblue eye in planet的區別?感覺這應該是 NLP 的領域了。

  1. stable diffusion 生成一個隨機的 latent space matrix。這個由 Random 決定,如果 Random 不變,則這個 latent space matrix 不變。

  2. 通過 noise-predictor,將 noisy image 與 text prompt 作爲入參,預測 predicted noise in latent space

  3. latent noise 減去 predicted noise。將其作爲新的 latent noise

  4. 不斷重複 2~3 執行 step 次。比如說 step=20

  5. 最終,通過 VAE 的 decoder 將 latent representation 生成最終的圖片

這個時候就可以貼 Stable diffusion 論文中的一張圖了

手撕一下公式:

左上角的定義爲一張 RGB 像素空間的圖。經過的變化,生成這個 latent space representation。再經過一系列的 noise encoder,得到,T 表示 step。

而這個過程則是 img2img 的 input。如果是 img2img,那麼初始的 noise latent representation 就是這個不斷加 noise 之後的

如果是 tex2img,初始的 noise latent representation 則是直接 random 出來的。

再從右下角的開始,y 表示多樣的控制條件的入參,如 text prompts。通過 (domain specific encoder) 將 y 轉爲 intermediate representation(一種中間表達)。而將經過 cross-attention layer 的實現:

具體的細節說實話沒看懂,而這一部分在 controlnet 中也有解釋,打算從 controlnet 的部分進行理解。

圖中 cross-attention 的部分可以很清晰的看到是一個由大到小,又由小到大的過程,在 controlnet 的圖中有解釋:

SD Encoder Block_1(64x64) -> SD Encoder Block_2(32x32) -> SD Encoder Block_3(16x16) -> SD Encoder(Block_4 8x8) -> SD Middle(Block 8x8) -> SD Decoder(Block_4 8x8) -> SD Decoder Block_3(16x16) -> SD Decoder Block_2(32x32) -> SD Decoder Blocker_1(64x64)

是一個從64x64 -> 8x8 -> 64x64的過程,具體爲啥,得等我撕完 controlnet 的論文。回到過程圖中,我們可以看到 denoising step 則是在 Latent Space 的左下角進行了一個循環,這裏與上面的流程一直。

最終通過 VAE 的 decoder D,輸出圖片

最終的公式如下:

結合上面的圖看,基本還是比較清晰的,不過這個:= 和代表了啥就不是很清楚了。結合 python 代碼看流程更清晰~ 刪掉了部分代碼,只留下了關鍵的調用。

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16
)
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
unet = UNet2DConditionModel.from_pretrained(
    "CompVis/stable-diffusion-v1-4", subfolder="unet"
)
scheduler = LMSDiscreteScheduler.from_pretrained(
    "CompVis/stable-diffusion-v1-4", subfolder="scheduler"
)
prompt = ["a photograph of an astronaut riding a horse"]
generator = torch.manual_seed(32)
text_input = tokenizer(
    prompt,
    padding="max_length",
    max_length=tokenizer.model_max_length,
    truncation=True,
    return_tensors="pt",
)
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
    uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
latents = torch.randn(
    (batch_size, unet.in_channels, height // 8, width // 8), generator=generator
)
scheduler.set_timesteps(num_inference_steps)
latents = latents * scheduler.init_noise_sigma
for t in tqdm(scheduler.timesteps):
    latent_model_input = torch.cat([latents] * 2)
    latent_model_input = scheduler.scale_model_input(latent_model_input, t)
    with torch.no_grad():
        noise_pred = unet(
            latent_model_input, t, encoder_hidden_states=text_embeddings
        ).sample
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (
        noise_pred_text - noise_pred_uncond
    )
    latents = scheduler.step(noise_pred, t, latents).prev_sample
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

還是很貼合圖中流程的。
在代碼中有一個 Scheduler,其實就是 noising 的執行器,它主要控制每一步 noising 的強度。
由 Scheduler 不斷加噪,然後 noise predictor 進行預測減噪。
具體可以看 Stable Diffusion Samplers: A Comprehensive Guide(地址:https://stable-diffusion-art.com/samplers/)

Img2Img

這個其實在上面的流程圖中已經解釋了。這裏把步驟列一下:

  1. 輸入的 image,通過 VAE 的 encoder 變成 latent space representation

  2. 往裏面加 noise,總共加 T 個 noise,noise 的強度由 Denoising strength 控制。noise 其實沒有循環加的過程,就是不斷疊同一個 noise T 次,所以可以一次計算完成。

  3. noisy image 和 text prompt 作爲輸入,由 noise predictor U-Net 預測一個新的 noise

  4. noisy image 減去預測的 noise

  5. 重複 3~4 step 次

  6. 通過 VAE 的 decoder 將 latent representation 轉變成 image

Inpainting

基於上面的原理,Inpainting 就很簡單了,noise 只加到 inpaint 的部分。其他和 Img2Img 一樣。相當於只生成 inpaint 的部分。所以我們也經常發現 inpaint 的邊緣經常無法非常平滑~ 如果能接受圖片的細微變化,可以調低 Denoising strength,將 inpaint 的結果,再進行一次 img2img

Stable Diffusion v1 vs v2

v2 開始 CLIP 的部分用了 OpenClip。導致生成的控制變得非常的難。OpenAI 的 CLIP 雖然訓練集更小,參數也更少。(OpenClip 是 ViT-L/14 CLIP 的 5 倍大小)。但似乎 ViT-L/14 的訓練集更好一些,有更多針對藝術和名人照片的部分,所以輸出的結果通常會更好。導致 v2 基本沒用起來。不過現在沒事了,SDXL 橫空出世。

SDXL model

SDXL 模型的參數達到了 66 億,而 v1.5 只有 9.8 億

由一個 Base model 和 Refiner model 組成。Base model 負責生成,而 Refiner 則負責加細節完善。可以只運行 Base model。但類似人臉眼睛模糊之類的問題還是需要 Refiner 解決。

SDXL 的主要變動:

  1. text encoder 組合了 OpenClip 和 ViT-G/14。畢竟 OpenClip 是可訓練的。

  2. 訓練用的圖片可以小於 256x256,增加了 39% 的訓練集

  3. U-Net 的部分比 v1.5 大了 3 倍

  4. 默認輸出就是 1024x1024

展示下對比效果:

從目前來看,有朝一日 SDXL 遲早替代 v1.5。從效果來說 v2.1 確實被時代淘汰了。

Stable Diffusion 的一些常見問題

臉部細節不足,比如說眼部模糊

可以通過 VAE files 進行修復~ 有點像 SDXL 的 Refiner

多指、少指

這個看起來是一個無解的問題。Andrew 給出的建議是加 prompt 比如說beautiful handsdetailed fingers,期望其中有部分圖片滿足要求。或者用 inpaint。反覆重新生成手部。(這個時候可以用相同的 prompt。)

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/0pdEorUcGCmSR02-DFRENA