使用🧨 Diffusers 進行 Stable Diffusion
Stable Diffusion 🎨 ...使用 🧨 Diffusers
Stable Diffusion 是一個文字到影像的潛在擴散模型,由 CompVis、Stability AI 和 LAION 的研究人員和工程師建立。它在 LAION-5B 資料庫的一個子集中訓練了 512x512 影像。LAION-5B 是目前最大、免費訪問的多模態資料集。
在這篇文章中,我們將展示如何使用 🧨 Diffusers 庫 來使用 Stable Diffusion,解釋模型的工作原理,並深入探討 diffusers
如何允許自定義影像生成管道。
注意:強烈建議對擴散模型的工作原理有基本的瞭解。如果您是擴散模型的新手,我們建議閱讀以下部落格文章之一
現在,讓我們開始生成一些影像 🎨。
執行 Stable Diffusion
許可證
在使用模型之前,您需要接受模型許可證才能下載和使用權重。注意:不再需要透過 UI 明確接受許可證。
該許可證旨在減輕如此強大的機器學習系統可能產生的有害影響。我們請求使用者完整而仔細地閱讀許可證。這裡我們提供一個摘要
- 您不能使用模型故意生成或分享非法或有害的輸出或內容,
- 我們不主張對您生成的輸出擁有任何權利,您可以自由使用它們,並對您的使用負責,您的使用不應違反許可證中規定的條款,以及
- 您可以重新分發權重並商業化和/或作為服務使用模型。如果您這樣做,請注意您必須包含與許可證中相同的限制,並將 CreativeML OpenRAIL-M 的副本分享給所有使用者。
用法
首先,您應該安裝 diffusers==0.10.2
來執行以下程式碼片段
pip install diffusers==0.10.2 transformers scipy ftfy accelerate
在本文中,我們將使用模型版本 v1-4
,但您也可以使用其他版本,例如 1.5、2 和 2.1,只需少量程式碼更改。
Stable Diffusion 模型只需幾行程式碼即可透過 StableDiffusionPipeline
管道進行推理。該管道設定了您從文字生成影像所需的一切,只需簡單呼叫 from_pretrained
函式。
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
如果有 GPU 可用,我們將其移到 GPU 上!
pipe.to("cuda")
注意:如果您的 GPU 記憶體有限,並且可用 GPU RAM 少於 10GB,請務必以 float16 精度(而不是預設的 float32 精度)載入 StableDiffusionPipeline
,如上所示。
您可以透過從 fp16
分支載入權重並告知 diffusers
預期權重為 float16 精度來完成此操作
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
要執行管道,只需定義提示並呼叫 pipe
即可。
prompt = "a photograph of an astronaut riding a horse"
image = pipe(prompt).images[0]
# you can save the image with
# image.save(f"astronaut_rides_horse.png")
結果如下所示
之前的程式碼每次執行時都會生成不同的影像。
如果您在某個時候獲得了黑色影像,這可能是因為模型中內建的內容過濾器檢測到 NSFW 結果。如果您認為不應該出現這種情況,請嘗試調整您的提示或使用不同的種子。實際上,模型預測中包含了有關特定結果是否檢測到 NSFW 的資訊。讓我們看看它們是什麼樣子
result = pipe(prompt)
print(result)
{
'images': [<PIL.Image.Image image mode=RGB size=512x512>],
'nsfw_content_detected': [False]
}
如果您想要確定性輸出,可以設定一個隨機種子並將生成器傳遞給管道。每次使用具有相同生成器的相同種子時,您都會得到相同的影像輸出。
import torch
generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]
# you can save the image with
# image.save(f"astronaut_rides_horse.png")
結果如下所示
您可以使用 num_inference_steps
引數更改推理步數。
一般來說,步數越多結果越好,但步數越多,生成時間越長。Stable Diffusion 在相對較少的步數下也能很好地工作,因此我們建議使用預設的 50
推理步數。如果您想要更快的結果,可以使用較小的步數。如果您想要可能更高質量的結果,可以使用較大的步數。
讓我們嘗試使用較少的去噪步驟執行管道。
import torch
generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, num_inference_steps=15, generator=generator).images[0]
# you can save the image with
# image.save(f"astronaut_rides_horse.png")
請注意,結構相同,但在宇航員的宇航服和馬的整體形態上存在問題。這表明僅使用 15 個去噪步驟已顯著降低了生成結果的質量。如前所述,通常 50
個去噪步驟足以生成高質量的影像。
除了 num_inference_steps
,我們在之前所有示例中都使用了另一個函式引數,稱為 guidance_scale
。guidance_scale
是一種提高對引導生成(在本例中是文字)的條件訊號的遵循程度以及整體樣本質量的方法。它也被稱為無分類器引導,簡單來說就是強制生成結果更好地匹配提示,這可能會以影像質量或多樣性為代價。對於 Stable Diffusion,介於 7
和 8.5
之間的值通常是不錯的選擇。預設情況下,管道使用 7.5 的 guidance_scale
。
如果您使用一個非常大的值,影像可能看起來不錯,但多樣性會降低。您可以在本文的此部分瞭解此引數的技術細節。
接下來,我們看看如何一次生成多個相同提示的影像。首先,我們將建立一個 image_grid
函式來幫助我們將其漂亮地視覺化為一個網格。
from PIL import Image
def image_grid(imgs, rows, cols):
assert len(imgs) == rows*cols
w, h = imgs[0].size
grid = Image.new('RGB', size=(cols*w, rows*h))
grid_w, grid_h = grid.size
for i, img in enumerate(imgs):
grid.paste(img, box=(i%cols*w, i//cols*h))
return grid
我們可以透過簡單地使用一個重複多次相同提示的列表來為同一個提示生成多張影像。我們將把列表傳送給管道,而不是之前使用的字串。
num_images = 3
prompt = ["a photograph of an astronaut riding a horse"] * num_images
images = pipe(prompt).images
grid = image_grid(images, rows=1, cols=3)
# you can save the grid with
# grid.save(f"astronaut_rides_horse.png")
預設情況下,Stable Diffusion 生成 512 × 512
畫素的影像。使用 height
和 width
引數很容易覆蓋預設值,以建立縱向或橫向比例的矩形影像。
選擇影像尺寸時,我們建議如下
- 確保
height
和width
都是8
的倍數。 - 低於 512 可能會導致影像質量下降。
- 在兩個方向上超過 512 將重複影像區域(全域性一致性會丟失)。
- 建立非方形影像的最佳方法是在一個維度上使用
512
,在另一個維度上使用更大的值。
讓我們執行一個例子
prompt = "a photograph of an astronaut riding a horse"
image = pipe(prompt, height=512, width=768).images[0]
# you can save the image with
# image.save(f"astronaut_rides_horse.png")
Stable Diffusion 工作原理
瞭解了 Stable Diffusion 可以產生的高質量影像後,讓我們嘗試更好地理解模型的功能。
Stable Diffusion 基於一種特殊型別的擴散模型,稱為潛在擴散,在《使用潛在擴散模型進行高解析度影像合成》中提出。
一般來說,擴散模型是機器學習系統,它們經過訓練,可以一步步地對隨機高斯噪聲進行去噪,以獲得感興趣的樣本,例如影像。有關它們如何工作的更詳細概述,請檢視此 colab。
擴散模型在生成影像資料方面已顯示出達到最先進水平的結果。但擴散模型的一個缺點是,由於其重複、順序的性質,反向去噪過程很慢。此外,這些模型會消耗大量記憶體,因為它們在畫素空間中執行,在生成高解析度影像時,畫素空間會變得非常大。因此,訓練這些模型和將其用於推理都具有挑戰性。
潛在擴散可以透過在低維潛在空間而不是實際畫素空間上應用擴散過程來降低記憶體和計算複雜性。這是標準擴散和潛在擴散模型之間的關鍵區別:在潛在擴散中,模型被訓練來生成影像的潛在(壓縮)表示。
潛在擴散主要有三個組成部分。
- 自編碼器(VAE)。
- U-Net。
- 一個文字編碼器,例如 CLIP 的文字編碼器。
1. 自編碼器 (VAE)
VAE 模型包含兩個部分,編碼器和解碼器。編碼器用於將影像轉換為低維潛在表示,該表示將作為 U-Net 模型的輸入。相反,解碼器將潛在表示轉換回影像。
在潛在擴散訓練過程中,編碼器用於獲取影像的潛在表示(潛變數)以進行前向擴散過程,該過程在每一步施加越來越多的噪聲。在推理過程中,反向擴散過程生成的去噪潛變數透過 VAE 解碼器轉換回影像。正如我們將在推理中看到的,我們只需要 VAE 解碼器。
2. U-Net
U-Net 有一個編碼器部分和一個解碼器部分,兩者都由 ResNet 塊組成。編碼器將影像表示壓縮為較低解析度的影像表示,解碼器將較低解析度的影像表示解碼回原始的較高解析度影像表示,據稱噪聲較少。更具體地說,U-Net 的輸出預測噪聲殘差,可用於計算預測的去噪影像表示。
為了防止 U-Net 在下采樣時丟失重要資訊,通常會在編碼器的下采樣 ResNet 和解碼器的上取樣 ResNet 之間新增快捷連線。此外,Stable Diffusion U-Net 能夠透過交叉注意力層以文字嵌入為條件輸出。交叉注意力層通常新增到 U-Net 的編碼器和解碼器部分,通常在 ResNet 塊之間。
3. 文字編碼器
文字編碼器負責將輸入提示(例如“一個騎馬的宇航員”)轉換為 U-Net 可以理解的嵌入空間。它通常是一個簡單的基於 Transformer 的編碼器,將輸入標記序列對映到潛在文字嵌入序列。
受Imagen的啟發,Stable Diffusion在訓練期間不訓練文字編碼器,而是簡單地使用CLIP的已訓練文字編碼器,即CLIPTextModel。
為什麼潛在擴散既快速又高效?
由於潛在擴散在低維空間上操作,與畫素空間擴散模型相比,它大大降低了記憶體和計算需求。例如,Stable Diffusion 中使用的自編碼器具有 8 倍的縮減因子。這意味著形狀為 (3, 512, 512)
的影像在潛在空間中變為 (4, 64, 64)
,這意味著空間壓縮比為 8 × 8 = 64
。
這就是為什麼即使在 16GB Colab GPU 上,也能如此快速地生成 512 × 512
影像!
推理過程中的 Stable Diffusion
綜上所述,現在讓我們透過說明邏輯流程來仔細研究模型如何在推理中工作。
Stable Diffusion 模型同時接收潛在種子和文字提示作為輸入。然後,潛在種子用於生成大小為 的隨機潛在影像表示,而文字提示則透過 CLIP 的文字編碼器轉換為大小為 的文字嵌入。
接下來,U-Net 在文字嵌入的條件下迭代地去噪隨機潛在影像表示。U-Net 的輸出,即噪聲殘差,用於透過排程器演算法計算去噪後的潛在影像表示。許多不同的排程器演算法可用於此計算,每種演算法都有其優點和缺點。對於 Stable Diffusion,我們建議使用以下演算法之一
排程器演算法函式的工作原理超出了本筆記本的範圍,但簡而言之,應該記住它們從先前的噪聲表示和預測的噪聲殘差計算預測的去噪影像表示。有關更多資訊,我們建議查閱《闡明基於擴散的生成模型的設計空間》
去噪過程重複約 50 次,以逐步檢索更好的潛在影像表示。完成後,潛在影像表示由變分自編碼器的解碼器部分解碼。
在對潛在和 Stable Diffusion 進行了簡要介紹之後,讓我們看看如何高階使用 🤗 Hugging Face diffusers
庫!
編寫自己的推理管道
最後,我們展示瞭如何使用 diffusers
建立自定義擴散管道。編寫自定義推理管道是 diffusers
庫的高階用法,可用於切換某些元件,例如上面解釋的 VAE 或排程器。
例如,我們將展示如何使用 Stable Diffusion 和不同的排程器,即Katherine Crowson的 K-LMS 排程器,該排程器在此 PR中新增。
預訓練模型包含設定完整擴散管道所需的所有元件。它們儲存在以下資料夾中
text_encoder
:Stable Diffusion 使用 CLIP,但其他擴散模型可能使用其他編碼器,例如BERT
。tokenizer
。它必須與text_encoder
模型使用的匹配。scheduler
:在訓練過程中逐步向影像新增噪聲的排程演算法。unet
:用於生成輸入潛在表示的模型。vae
:我們將用於將潛在表示解碼為真實影像的自編碼器模組。
我們可以透過引用它們儲存的資料夾,使用 from_pretrained
的 subfolder
引數來載入元件。
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
# 1. Load the autoencoder model which will be used to decode the latents into image space.
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
# 2. Load the tokenizer and text encoder to tokenize and encode the text.
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
# 3. The UNet model for generating the latents.
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
現在,我們不是載入預定義的排程器,而是載入具有一些合適引數的 K-LMS 排程器。
from diffusers import LMSDiscreteScheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
接下來,讓我們將模型移到 GPU 上。
torch_device = "cuda"
vae.to(torch_device)
text_encoder.to(torch_device)
unet.to(torch_device)
我們現在定義將用於生成影像的引數。
請注意,guidance_scale
的定義類似於Imagen 論文中公式 (2) 的引導權重 w
。guidance_scale == 1
對應於不進行無分類器引導。這裡我們將其設定為 7.5,如前所述。
與前面的示例不同,我們將 num_inference_steps
設定為 100,以獲得更清晰的影像。
prompt = ["a photograph of an astronaut riding a horse"]
height = 512 # default height of Stable Diffusion
width = 512 # default width of Stable Diffusion
num_inference_steps = 100 # Number of denoising steps
guidance_scale = 7.5 # Scale for classifier-free guidance
generator = torch.manual_seed(0) # Seed generator to create the initial latent noise
batch_size = len(prompt)
首先,我們獲取傳遞的提示的 text_embeddings
。這些嵌入將用於條件化 UNet 模型,並引導影像生成接近輸入提示的內容。
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(torch_device))[0]
我們還將獲取用於無分類器引導的無條件文字嵌入,它們只是填充標記(空文字)的嵌入。它們需要與條件 text_embeddings
具有相同的形狀(batch_size
和 seq_length
)
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
[""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
對於無分類器引導,我們需要進行兩次前向傳播:一次使用條件輸入 (text_embeddings
),另一次使用無條件嵌入 (uncond_embeddings
)。在實踐中,我們可以將兩者連線成一個批次,以避免進行兩次前向傳播。
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
接下來,我們生成初始隨機噪聲。
latents = torch.randn(
(batch_size, unet.in_channels, height // 8, width // 8),
generator=generator,
)
latents = latents.to(torch_device)
如果我們此時檢查 latents
,我們會看到它們的形狀是 torch.Size([1, 4, 64, 64])
,比我們想要生成的影像小得多。模型稍後會將此潛在表示(純噪聲)轉換為 512 × 512
影像。
接下來,我們使用我們選擇的 num_inference_steps
初始化排程器。這將計算在去噪過程中使用的 sigmas
和精確時間步長值。
scheduler.set_timesteps(num_inference_steps)
K-LMS 排程器需要將其 latents
乘以其 sigma
值。讓我們在這裡完成它
latents = latents * scheduler.init_noise_sigma
我們已經準備好編寫去噪迴圈。
from tqdm.auto import tqdm
scheduler.set_timesteps(num_inference_steps)
for t in tqdm(scheduler.timesteps):
# expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
latent_model_input = torch.cat([latents] * 2)
latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)
# predict the noise residual
with torch.no_grad():
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
# perform guidance
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(noise_pred, t, latents).prev_sample
現在我們使用 vae
將生成的 latents
解碼回影像。
# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
image = vae.decode(latents).sample
最後,讓我們將影像轉換為 PIL 格式,以便我們可以顯示或儲存它。
image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images[0]
我們從使用 🤗 Hugging Face Diffusers 的 Stable Diffusion 的基本用法,到更高階的庫用法,並嘗試介紹了現代擴散系統中的所有部分。如果您喜歡這個主題並想了解更多,我們推薦以下資源
- 我們的 Colab 筆記本。
- 《Diffusers 入門》筆記本,它對擴散系統進行了更廣泛的概述。
- 《註釋擴散模型》部落格文章。
- 我們在 GitHub 上的程式碼,如果
diffusers
對您有用,我們非常樂意您點亮 ⭐!
引用:
@article{patil2022stable,
author = {Patil, Suraj and Cuenca, Pedro and Lambert, Nathan and von Platen, Patrick},
title = {Stable Diffusion with 🧨 Diffusers},
journal = {Hugging Face Blog},
year = {2022},
note = {[https://huggingface.co/blog/rlhf](https://huggingface.co/blog/stable_diffusion)},
}