【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）
  
  
上一篇博文我们学习了Stable Diffusion的原理，这一篇我们继续深入了解Stable Diffusion的应用实践和代码分析。
1 AutoEncoder

SD采用基于KL-reg的autoencoder，当输入图像为512x512时将得到64x64x4大小的latent。autoencoder模型是在OpenImages数据集上基于256x256大小训练的，但是由于模型是全卷积结构的（基于ResnetBlock），所以可以扩展应用在尺寸>256的图像上。
下面我们利用diffusers库来加载autoencoder模型，实现图像的压缩和重修，代码如下：

1. import torch  
2. from diffusers import AutoencoderKL  
3. import numpy as np  
4. from PIL import Image
6. print(torch.cuda.is_available())
7.   
8. #加载模型: autoencoder可以通过SD权重指定subfolder来单独加载  
9. print("Start...")
10. autoencoder = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
11. autoencoder.to("cuda", dtype=torch.float16)  
12. print("Get weight successfully")
13.   
14. # 读取图像并预处理  
15. # raw_image = Image.open("liuyifei.jpg").convert("RGB").resize((256, 256))  
16. raw_image = Image.open("liuyifei.jpg").convert("RGB")
17. image = np.array(raw_image).astype(np.float32) / 127.5 - 1.0  
18. image = image[None].transpose(0, 3, 1, 2)  
19. image = torch.from_numpy(image)  
20.   
21. # 压缩图像为latent并重建  
22. with torch.inference_mode():  
23.     latent = autoencoder.encode(image.to("cuda", dtype=torch.float16)).latent_dist.sample()  
24.     rec_image = autoencoder.decode(latent).sample  
25.     rec_image = (rec_image / 2 + 0.5).clamp(0, 1)  
26.     rec_image = rec_image.cpu().permute(0, 2, 3, 1).numpy()  
27.     rec_image = (rec_image * 255).round().astype("uint8")  
28.     rec_image = Image.fromarray(rec_image[0])  
30. rec_image.save("liuyifei_re.jpg")

复制代码

重修结果如下所示，对比手表上的文字，可以看出，autoencoder将图片压缩到latent后再重修其实是有损的。

为了改善这种畸变，stabilityai在发布SD 2.0时同时发布了两个在LAION子数据集上精调的autoencoder，注意这里只精调autoencoder的decoder部分，SD的UNet在训练过程只需要encoder部分，所以这样精调后的autoencoder可以直接用在先前训练好的UNet上（这种技巧照旧比力通用的，好比谷歌的Parti也是在训练好后自回归生成模型后，扩大并精调ViT-VQGAN的decoder模块来提升生成质量）。我们也可以直接在diffusers中利用这些autoencoder，好比mse版本（采用mse丧失来finetune的模型）：

1. 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的代码如下：

1. from transformers import CLIPTextModel, CLIPTokenizer  
2.   
3. text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder").to("cuda")  
4. # text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")  
5. tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")  
6. # tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")  
7.   
8. # 对输入的text进行tokenize，得到对应的token ids  
9. prompt = "a photograph of an astronaut riding a horse"  
10. text_input_ids = tokenizer(  
11.     prompt,  
12.     padding="max_length",  
13.     max_length=tokenizer.model_max_length,  
14.     truncation=True,  
15.     return_tensors="pt"  
16. ).input_ids
18. print(f" \n\n    text_input_ids: {text_input_ids}   \n\n")
19.   
20. # 将token ids送入text model得到77x768的特征  
21. text_embeddings = text_encoder(text_input_ids.to("cuda"))[0]  
22. print(f" \n\n    text_embeddings: {text_embeddings}   \n\n")

复制代码

输出如下：

1.     text_input_ids: tensor([[49406,   320,  8853,   539,   550, 18376,  6765,   320,  4558, 49407,
2.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
3.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
4.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
5.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
6.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
7.          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
8.          49407, 49407, 49407, 49407, 49407, 49407, 49407]])
10.     text_embeddings: tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
11.          [ 0.0290, -1.3258,  0.3085,  ..., -0.5257,  0.9768,  0.6652],
12.          [ 0.4595,  0.5617,  1.6663,  ..., -1.9515, -1.2307,  0.0104],
13.          ...,
14.          [-3.0421, -0.0656, -0.1793,  ...,  0.3943, -0.0190,  0.7664],
15.          [-3.0551, -0.1036, -0.1936,  ...,  0.4236, -0.0189,  0.7575],
16.          [-2.9854, -0.0832, -0.1715,  ...,  0.4355,  0.0095,  0.7485]]],
17.        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模块是不包罗下采样的。

其中CrossAttnDownBlock2D模块的主要结构如下图所示，text condition将通过CrossAttention模块嵌入进来，此时Attention的query是UNet的中间特征，而key和value则是text embeddings。

SD和DDPM一样采用猜测noise的方法来训练UNet，其训练丧失也和DDPM一样。基于diffusers库，我们可以实现SD的训练，其核心代码如下：

1. import torch
2. from torch.utils.data import Dataset, DataLoader
3. from torchvision import transforms
4. from PIL import Image
5. from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
6. from transformers import CLIPTextModel, CLIPTokenizer
7. import torch.nn.functional as F
9. # 自定义Dataset类
10. class CustomImageTextDataset(Dataset):
11.     def __init__(self, image_paths, text_descriptions, transform=None):
12.         self.image_paths = image_paths
13.         self.text_descriptions = text_descriptions
14.         self.transform = transform
15.         
16.     def __len__(self):
17.         return len(self.image_paths)
18.    
19.     def __getitem__(self, idx):
20.         image_path = self.image_paths[idx]
21.         text_description = self.text_descriptions[idx]
22.         
23.         # 加载图像
24.         image = Image.open(image_path).convert("RGB")
25.         if self.transform:
26.             image = self.transform(image)
27.         
28.         return {
29.             'image': image,
30.             'text': text_description
31.         }
33. # 数据准备
34. image_paths = ["path/to/image1.jpg", "path/to/image2.jpg"]  # 替换为实际的图像路径
35. text_descriptions = ["description for image1", "description for image2"]  # 替换为实际的文本描述
37. # 图像转换（预处理）
38. transform = transforms.Compose([
39.     transforms.Resize((256, 256)),  # 调整大小
40.     transforms.ToTensor(),  # 转换为张量
41. ])
43. # 创建数据集实例
44. dataset = CustomImageTextDataset(image_paths=image_paths, text_descriptions=text_descriptions, transform=transform)
46. # 创建DataLoader
47. train_dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)
49. # 加载autoencoder
50. vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
51. # 加载text encoder
52. text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
53. tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
55. model_config = {
56.     "sample_size": 32,
57.     "in_channels": 4,
58.     "out_channels": 4,
59.     "down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"),
60.     "up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
61.     "block_out_channels": (320, 640, 1280, 1280),
62.     "layers_per_block": 2,
63.     "cross_attention_dim": 768,
64.     "attention_head_dim": 8,
65. }
66. # 初始化UNet
67. unet = UNet2DConditionModel(**model_config)
69. # 定义scheduler
70. noise_scheduler = DDPMScheduler(
71.     beta_start=0.00085,
72.     beta_end=0.012,
73.     beta_schedule="scaled_linear",
74.     num_train_timesteps=1000
75. )
77. # 冻结vae和text_encoder
78. vae.requires_grad_(False)
79. text_encoder.requires_grad_(False)
81. opt = torch.optim.AdamW(unet.parameters(), lr=1e-4)
83. # 训练循环
84. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
85. unet.to(device)
86. vae.to(device)
87. text_encoder.to(device)
89. for epoch in range(10):  # 假设训练10个epoch
90.     unet.train()
91.     for step, batch in enumerate(train_dataloader):
92.         with torch.no_grad():
93.             # 将image转到latent空间
94.             latents = vae.encode(batch["image"].to(device)).latent_dist.sample()
95.             # rescaling latents
96.             latents = latents * vae.config.scaling_factor
97.             
98.             # 提取text embeddings
99.             text_input_ids = tokenizer(
100.                 batch["text"],
101.                 padding="max_length",
102.                 max_length=tokenizer.model_max_length,
103.                 truncation=True,
104.                 return_tensors="pt"
105.             ).input_ids.to(device)
106.             text_embeddings = text_encoder(text_input_ids)[0]
108.         # 随机采样噪音
109.         noise = torch.randn_like(latents)
110.         bsz = latents.shape[0]
112.         # 随机采样timestep
113.         timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=device).long()
115.         # 将noise添加到latent上，即扩散过程
116.         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
118.         # 预测noise并计算loss
119.         model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
120.         loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
122.         opt.zero_grad()
123.         loss.backward()
124.         opt.step()
126.         if step % 10 == 0:
127.             print(f"Epoch {epoch}, Step {step}, Loss: {loss.item()}")
129. # 在训练完成后保存模型
130. model_save_path = 'path/to/your/unet_model.pth'
131. torch.save(unet.state_dict(), model_save_path)
132. print(f"Model has been saved to {model_save_path}")
134. optimizer_save_path = 'path/to/your/optimizer.pth'
135. torch.save(opt.state_dict(), optimizer_save_path)
136. print(f"Optimizer state has been saved to {optimizer_save_path}")
138. # 加载模型进行推理或继续训练
139. unet_load_path = 'path/to/your/unet_model.pth'
140. unet_loaded = UNet2DConditionModel(**model_config)  # 创建一个与原模型结构相同的实例
141. unet_loaded.load_state_dict(torch.load(unet_load_path))
142. unet_loaded.to(device)
143. unet_loaded.eval()  # 设置为评估模式
145. # 恢复优化器的状态以继续训练
146. opt_load_path = 'path/to/your/optimizer.pth'
147. opt_loaded = torch.optim.AdamW(unet_loaded.parameters(), lr=1e-4)  # 创建一个新的优化器实例
148. opt_loaded.load_state_dict(torch.load(opt_load_path))
150. # 使用unet_loaded进行推理或者用opt_loaded继续训练。

复制代码

注意的是SD的noise scheduler虽然也是采用一个1000步长的scheduler，但是不是linear的，而是scaled linear，具体的计算如下所示：

1. 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来实现文生图，具体代码如下所示：

1. import torch  
2. from diffusers import StableDiffusionPipeline  
3. from PIL import Image  
5. # 组合图像，生成grid  
6. def image_grid(imgs, rows, cols):  
7.     assert len(imgs) == rows*cols  
9.     w, h = imgs[0].size  
10.     grid = Image.new('RGB', size=(cols*w, rows*h))  
11.     grid_w, grid_h = grid.size  
13.     for i, img in enumerate(imgs):  
14.         grid.paste(img, box=(i%cols*w, i//cols*h))  
15.     return grid  
17. # 加载文生图pipeline  
18. pipe = StableDiffusionPipeline.from_pretrained(  
19.     "runwayml/stable-diffusion-v1-5", # 或者使用 SD v1.4: "CompVis/stable-diffusion-v1-4"  
20.     torch_dtype=torch.float16  
21. ).to("cuda")  
23. # 输入text，这里text又称为prompt  
24. prompts = [  
25.     "a photograph of an astronaut riding a horse",  
26.     "A cute otter in a rainbow whirlpool holding shells, watercolor",  
27.     "An avocado armchair",  
28.     "A white dog wearing sunglasses"  
29. ]  
31. generator = torch.Generator("cuda").manual_seed(42) # 定义随机seed，保证可重复性  
33. # 执行推理  
34. images = pipe(  
35.     prompts,  
36.     height=512,  
37.     width=512,  
38.     num_inference_steps=50,  
39.     guidance_scale=7.5,  
40.     negative_prompt=None,  
41.     num_images_per_prompt=1,  
42.     generator=generator  
43. ).images  
45. # 保存每个单独的图片
46. for idx, img in enumerate(images):
47.     img.save(f"image_{idx}.png")
49. # 创建并保存组合后的网格图
50. grid = image_grid(images, rows=1, cols=len(prompts))
51. grid.save("combined_images.png")
52. 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：
  

1. from diffusers import DDIMScheduler  
2.   
3. # 注意这里的clip_sample要关闭，否则生成图像存在问题，因为不能对latent进行clip  
4. 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来实现文生图，具体代码如下所示：

1. import torch
2. from diffusers import StableDiffusionImg2ImgPipeline
3. from PIL import Image
5. # 加载图生图pipeline
6. model_id = "runwayml/stable-diffusion-v1-5"
7. pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
9. # 读取初始图片
10. init_image = Image.open("liuyifei.jpg").convert("RGB").resize((512, 512))
11. print(init_image.size)
12. init_image.save("liuyifei_512.jpg")
14. # 推理
15. prompt = "A girl wearing a hat on her head."
16. generator = torch.Generator(device="cuda").manual_seed(2023)
18. image = pipe(
19.     prompt=prompt,
20.     image=init_image,
21.     strength=0.8,
22.     guidance_scale=7.5,
23.     generator=generator
24. ).images[0]
26. # 保存生成的图像
27. output_path = "generated_liuyifei.jpg"
28. image.save(output_path)
29. 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，具体代码如下所示：

1. import torch  
2. from diffusers import StableDiffusionInpaintPipelineLegacy  
3. from PIL import Image  
4.   
5. # 加载inpainting pipeline  
6. model_id = "runwayml/stable-diffusion-v1-5"  
7. pipe = StableDiffusionInpaintPipelineLegacy.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")  
8.   
9. # 读取输入图像和输入mask  
10. input_image = Image.open("overture-creations-5sI6fQgYIuo.png").resize((512, 512))  
11. input_mask = Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512, 512))  
12.   
13. # 执行推理  
14. prompt = ["a mecha robot sitting on a bench", "a cat sitting on a bench"]  
15. generator = torch.Generator("cuda").manual_seed(0)  
16.   
17. with torch.autocast("cuda"):  
18.     images = pipe(  
19.         prompt=prompt,  
20.         image=input_image,  
21.         mask_image=input_mask,  
22.         num_inference_steps=50,  
23.         strength=0.75,  
24.         guidance_scale=7.5,  
25.         num_images_per_prompt=1,  
26.         generator=generator,  
27.     ).images  
30. # 保存每个单独的图片
31. for idx, img in enumerate(images):
32.     img.save(f"image_{idx}.png")
34. print("所有图片已保存到本地。")

复制代码

5 其他

Colab上开源的Stable Diffusion 2.1 GUI：stable_diffusion_2_0.ipynb。
最强盛且模块化的具有图形/节点界面的稳定扩散GUI：ComfyUI。
Huggingface模型库：https://huggingface.co/stabilityai。
Huggingface的Diffuser库：https://github.com/huggingface/diffusers。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。