【AIGC系列】4：Stable Diffusion應用實踐和代碼分析

AIGC系列博文：
【AIGC系列】1：自編碼器（AutoEncoder, AE）
【AIGC系列】2：DALL·E 2模型介紹（內含擴散模型介紹）
【AIGC系列】3：Stable Diffusion模型原理介紹
【AIGC系列】4：Stable Diffusion應用實踐和代碼分析
【AIGC系列】5：視頻生成模型數據處理和預訓練流程介紹（Sora、MovieGen、HunyuanVideo）

1 AutoEncoder

SD采用基于KL-reg的autoencoder，當輸入圖像為512x512時將得到64x64x4大小的latent。autoencoder模型是在OpenImages數據集上基于256x256大小訓練的，但是由于模型是全卷積結構的（基于ResnetBlock），所以可以擴展應用在尺寸>256的圖像上。

下面我們使用diffusers庫來加載autoencoder模型，實現圖像的壓縮和重建，代碼如下：

import torch  
from diffusers import AutoencoderKL  
import numpy as np  
from PIL import Imageprint(torch.cuda.is_available())#加載模型: autoencoder可以通過SD權重指定subfolder來單獨加載  
print("Start...")
autoencoder = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
autoencoder.to("cuda", dtype=torch.float16)  
print("Get weight successfully")# 讀取圖像并預處理  
# raw_image = Image.open("liuyifei.jpg").convert("RGB").resize((256, 256))  
raw_image = Image.open("liuyifei.jpg").convert("RGB")
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.save("liuyifei_re.jpg")

重建效果如下所示，對比手表上的文字，可以看出，autoencoder將圖片壓縮到latent后再重建其實是有損的。

對比1

為了改善這種畸變，stabilityai在發布SD 2.0時同時發布了兩個在LAION子數據集上精調的autoencoder，注意這里只精調autoencoder的decoder部分，SD的UNet在訓練過程只需要encoder部分，所以這樣精調后的autoencoder可以直接用在先前訓練好的UNet上（這種技巧還是比較通用的，比如谷歌的Parti也是在訓練好后自回歸生成模型后，擴大并精調ViT-VQGAN的decoder模塊來提升生成質量）。我們也可以直接在diffusers中使用這些autoencoder，比如mse版本（采用mse損失來finetune的模型）：

autoencoder = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse/")

2 CLIP text encoder

SD采用CLIP text encoder來對輸入的文本生成text embeddings，采用的CLIP模型是clip-vit-large-patch14，該模型的text encoder層數為12，特征維度為768，模型參數大小是123M。文本輸入text encoder后得到最后的hidden states特征維度大小為77x768（77是token的數量），這個細粒度的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 = tokenizer(  prompt,  padding="max_length",  max_length=tokenizer.model_max_length,  truncation=True,  return_tensors="pt"  
).input_idsprint(f" \n\n    text_input_ids: {text_input_ids}   \n\n")# 將token ids送入text model得到77x768的特征  
text_embeddings = text_encoder(text_input_ids.to("cuda"))[0]  
print(f" \n\n    text_embeddings: {text_embeddings}   \n\n")

輸出如下：

    text_input_ids: tensor([[49406,   320,  8853,   539,   550, 18376,  6765,   320,  4558, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,49407, 49407, 49407, 49407, 49407, 49407, 49407]])text_embeddings: tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],[ 0.0290, -1.3258,  0.3085,  ..., -0.5257,  0.9768,  0.6652],[ 0.4595,  0.5617,  1.6663,  ..., -1.9515, -1.2307,  0.0104],...,[-3.0421, -0.0656, -0.1793,  ...,  0.3943, -0.0190,  0.7664],[-3.0551, -0.1036, -0.1936,  ...,  0.4236, -0.0189,  0.7575],[-2.9854, -0.0832, -0.1715,  ...,  0.4355,  0.0095,  0.7485]]],device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

值得注意的是，這里的tokenizer最大長度為77（CLIP訓練時所采用的設置），當輸入text的tokens數量超過77后，將進行截斷，如果不足則進行paddings，這樣將保證無論輸入任何長度的文本（甚至是空文本）都得到77x768大小的特征。在上面的例子里，輸入的tokens數量少于77，所以后面都padding了id為49407的token。

在訓練SD的過程中，CLIP text encoder模型是凍結的。在早期的工作中，比如OpenAI的GLIDE和latent diffusion中的LDM均采用一個隨機初始化的tranformer模型來提取text的特征，但是最新的工作都是采用預訓練好的text model。比如谷歌的Imagen采用純文本模型T5 encoder來提出文本特征，而SD則采用CLIP text encoder，預訓練好的模型往往已經在大規模數據集上進行了訓練，它們要比直接采用一個從零訓練好的模型要好。

3 UNet

SD的擴散模型是一個860M的UNet，其主要結構如下圖所示，其中encoder部分包括3個CrossAttnDownBlock2D模塊和1個DownBlock2D模塊，而decoder部分包括1個UpBlock2D模塊和3個CrossAttnUpBlock2D模塊，中間還有一個UNetMidBlock2DCrossAttn模塊。

encoder和decoder兩個部分是完全對應的，中間有skip connection。3個CrossAttnDownBlock2D模塊最后均有一個2x的downsample操作，而DownBlock2D模塊是不包含下采樣的。

UNet

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

SD和DDPM一樣采用預測noise的方法來訓練UNet，其訓練損失也和DDPM一樣。基于diffusers庫，我們可以實現SD的訓練，其核心代碼如下：

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
import torch.nn.functional as F# 自定義Dataset類
class CustomImageTextDataset(Dataset):def __init__(self, image_paths, text_descriptions, transform=None):self.image_paths = image_pathsself.text_descriptions = text_descriptionsself.transform = transformdef __len__(self):return len(self.image_paths)def __getitem__(self, idx):image_path = self.image_paths[idx]text_description = self.text_descriptions[idx]# 加載圖像image = Image.open(image_path).convert("RGB")if self.transform:image = self.transform(image)return {'image': image,'text': text_description}# 數據準備
image_paths = ["path/to/image1.jpg", "path/to/image2.jpg"]  # 替換為實際的圖像路徑
text_descriptions = ["description for image1", "description for image2"]  # 替換為實際的文本描述# 圖像轉換（預處理）
transform = transforms.Compose([transforms.Resize((256, 256)),  # 調整大小transforms.ToTensor(),  # 轉換為張量
])# 創建數據集實例
dataset = CustomImageTextDataset(image_paths=image_paths, text_descriptions=text_descriptions, transform=transform)# 創建DataLoader
train_dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)# 加載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")model_config = {"sample_size": 32,"in_channels": 4,"out_channels": 4,"down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"),"up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),"block_out_channels": (320, 640, 1280, 1280),"layers_per_block": 2,"cross_attention_dim": 768,"attention_head_dim": 8,
}
# 初始化UNet
unet = UNet2DConditionModel(**model_config)# 定義scheduler
noise_scheduler = DDPMScheduler(beta_start=0.00085,beta_end=0.012,beta_schedule="scaled_linear",num_train_timesteps=1000
)# 凍結vae和text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)opt = torch.optim.AdamW(unet.parameters(), lr=1e-4)# 訓練循環
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
unet.to(device)
vae.to(device)
text_encoder.to(device)for epoch in range(10):  # 假設訓練10個epochunet.train()for step, batch in enumerate(train_dataloader):with torch.no_grad():# 將image轉到latent空間latents = vae.encode(batch["image"].to(device)).latent_dist.sample()# rescaling latentslatents = latents * vae.config.scaling_factor# 提取text embeddingstext_input_ids = tokenizer(batch["text"],padding="max_length",max_length=tokenizer.model_max_length,truncation=True,return_tensors="pt").input_ids.to(device)text_embeddings = text_encoder(text_input_ids)[0]# 隨機采樣噪音noise = torch.randn_like(latents)bsz = latents.shape[0]# 隨機采樣timesteptimesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=device).long()# 將noise添加到latent上，即擴散過程noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)# 預測noise并計算lossmodel_pred = unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sampleloss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")opt.zero_grad()loss.backward()opt.step()if step % 10 == 0:print(f"Epoch {epoch}, Step {step}, Loss: {loss.item()}")# 在訓練完成后保存模型
model_save_path = 'path/to/your/unet_model.pth'
torch.save(unet.state_dict(), model_save_path)
print(f"Model has been saved to {model_save_path}")optimizer_save_path = 'path/to/your/optimizer.pth'
torch.save(opt.state_dict(), optimizer_save_path)
print(f"Optimizer state has been saved to {optimizer_save_path}")# 加載模型進行推理或繼續訓練
unet_load_path = 'path/to/your/unet_model.pth'
unet_loaded = UNet2DConditionModel(**model_config)  # 創建一個與原模型結構相同的實例
unet_loaded.load_state_dict(torch.load(unet_load_path))
unet_loaded.to(device)
unet_loaded.eval()  # 設置為評估模式# 恢復優化器的狀態以繼續訓練
opt_load_path = 'path/to/your/optimizer.pth'
opt_loaded = torch.optim.AdamW(unet_loaded.parameters(), lr=1e-4)  # 創建一個新的優化器實例
opt_loaded.load_state_dict(torch.load(opt_load_path))# 使用unet_loaded進行推理或者用opt_loaded繼續訓練。

注意的是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

在訓練條件擴散模型時，往往會采用Classifier-Free Guidance （CFG），即在訓練條件擴散模型的同時也訓練一個無條件的擴散模型，同時在采樣階段將條件控制下預測的噪音和無條件下的預測噪音組合在一起來確定最終的噪音，CFG對于提升條件擴散模型的圖像生成效果是至關重要的。

4 應用

4.1 文生圖

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

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

import torch  
from diffusers import StableDiffusionPipeline  
from PIL import Image  # 組合圖像，生成grid  
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  # 加載文生圖pipeline  
pipe = StableDiffusionPipeline.from_pretrained(  "runwayml/stable-diffusion-v1-5", # 或者使用 SD v1.4: "CompVis/stable-diffusion-v1-4"  torch_dtype=torch.float16  
).to("cuda")  # 輸入text，這里text又稱為prompt  
prompts = [  "a photograph of an astronaut riding a horse",  "A cute otter in a rainbow whirlpool holding shells, watercolor",  "An avocado armchair",  "A white dog wearing sunglasses"  
]  generator = torch.Generator("cuda").manual_seed(42) # 定義隨機seed，保證可重復性  # 執行推理  
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  # 保存每個單獨的圖片
for idx, img in enumerate(images):img.save(f"image_{idx}.png")# 創建并保存組合后的網格圖
grid = image_grid(images, rows=1, cols=len(prompts))
grid.save("combined_images.png")
print("所有圖片已保存到本地。")

生成的結果如下：

天馬行空

重要參數說明：

指定width和height來決定生成圖像的大小：前面說過SD最后是在512x512尺度上訓練的，所以生成512x512尺寸效果是最好的，但是實際上SD可以生成任意尺寸的圖片：一方面autoencoder支持任意尺寸的圖片的編碼和解碼，另外一方面擴散模型UNet也是支持任意尺寸的latents生成的（UNet是卷積+attention的混合結構）。但是生成512x512以外的圖片會存在一些問題，比如生成低分辨率圖像時，圖像的質量大幅度下降等等。
num_inference_steps：指推理過程中的去噪步數或者采樣步數。SD在訓練過程采用的是步數為1000的noise scheduler，但是在推理時往往采用速度更快的scheduler：只需要少量的采樣步數就能生成不錯的圖像，比如SD默認采用PNDM scheduler，它只需要采樣50步就可以出圖。當然我們也可以換用其它類型的scheduler，比如DDIM scheduler和DPM-Solver scheduler。我們可以在diffusers中直接替換scheduler，比如我們想使用DDIM：

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

guidance_scale：當CFG的guidance_scale越大時，生成的圖像應該會和輸入文本更一致。SD默認采用的guidance_scale為7.5。但是過大的guidance_scale也會出現問題，主要是由于訓練和測試的不一致，過大的guidance_scale會導致生成的樣本超出范圍。
negative_prompt：這個參數和CFG有關，去噪過程的噪音預測不僅僅依賴條件擴散模型，也依賴無條件擴散模型，這里的negative_prompt便是無條件擴散模型的text輸入，前面說過訓練過程中我們將text置為空字符串來實現無條件擴散模型，所以這里negative_prompt = None 。但是有時候我們可以使用不為空的negative_prompt來避免模型生成的圖像包含不想要的東西，因為從上述公式可以看到這里的無條件擴散模型是我們想遠離的部分。

4.2 圖生圖

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

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

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

import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image# 加載圖生圖pipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")# 讀取初始圖片
init_image = Image.open("liuyifei.jpg").convert("RGB").resize((512, 512))
print(init_image.size)
init_image.save("liuyifei_512.jpg")# 推理
prompt = "A girl wearing a hat on her head."
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]# 保存生成的圖像
output_path = "generated_liuyifei.jpg"
image.save(output_path)
print(f"Generated image saved to {output_path}")

原始圖片：

效果如下：

相比文生圖的pipeline，圖生圖的pipeline還多了一個參數strength，這個參數介于0-1之間，表示對輸入圖片加噪音的程度，這個值越大加的噪音越多，對原始圖片的破壞也就越大，當strength=1時，其實就變成了一個隨機噪音，此時就相當于純粹的文生圖pipeline了。

4.3 圖像inpainting

圖像inpainting和圖生圖一樣也是文生圖功能的一個擴展。SD的圖像inpainting不是用在圖像修復上，而是主要用在圖像編輯上：給定一個輸入圖像和想要編輯的區域mask，我們想通過文生圖來編輯mask區域的內容。

它和圖生圖一樣，首先將輸入圖像通過autoencoder編碼為latent，然后加入一定的高斯噪音生成noisy latent，再進行去噪生成圖像，但是這里為了保證mask以外的區域不發生變化，在去噪過程的每一步，都將擴散模型預測的noisy latent用真實圖像同level的nosiy latent替換。

在diffusers中，使用StableDiffusionInpaintPipelineLegacy可以實現文本引導下的圖像inpainting，具體代碼如下所示：

import torch  
from diffusers import StableDiffusionInpaintPipelineLegacy  
from PIL import Image  # 加載inpainting pipeline  
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))  # 執行推理  
prompt = ["a mecha robot sitting on a bench", "a cat sitting on a bench"]  
generator = torch.Generator("cuda").manual_seed(0)  with torch.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  # 保存每個單獨的圖片
for idx, img in enumerate(images):img.save(f"image_{idx}.png")print("所有圖片已保存到本地。")