【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模型，实现图像的压缩和重建，代码如下：

import torch  
from diffusers import AutoencoderKL  
import numpy as np  
from PIL import Image

print(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后再重建其实是有损的。

为了改善这种畸变，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_ids

print(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模块是不包含下采样的。

其中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_paths
        self.text_descriptions = text_descriptions
        self.transform = transform
        
    def __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个epoch
    unet.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 latents
            latents = latents * vae.config.scaling_factor
            
            # 提取text embeddings
            text_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]

        # 随机采样timestep
        timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=device).long()

        # 将noise添加到latent上，即扩散过程
        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

        # 预测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.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("所有图片已保存到本地。")

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。