当前位置：首页 > article >正文

DIFFUSION 系列笔记| Latent Diffusion Model、Stable Diffusion基础概念、数学原理、代码分析、案例展示

article 2024/9/24 3:25:06

Latent Diffusion Model

LDM 主要思想

LDM使用示例

LDM Pipeline

LDM 中的 UNET

准备时间步 time steps

预处理阶段 pre-process

下采样过程 down sampling

中间处理 mid processing

上采样 upsampling

后处理 post-process

LDM Super Resolution Pipeline

Stable diffusion

SD v1 架构

1.Text Encoder

2.Diffusion 反向采样过程

3.Super resolution

SD v1.1 - v1.5

SD v2

Lora

Latent Diffusion Model

论文：High-Resolution Image Synthesis with Latent Diffusion Models

LDM 主要思想

扩散模型（DMs）直接在像素领域工作，优化和推断都很费时。为了在有限的计算资源上训练它们，LDM 先使用一个预训练好的 AutoEncoder，将图片像素转换到了维度较小的 latent space 上，而后再进行传统的扩散模型推理与优化。这种训练方式使得 LDM 在算力和性能之间得到了平衡。

此外，通过引入交叉注意力，使得 DMs 能够在条件生成上有不错的效果，包括如文字生成图片，inpainting 等。

LDM使用示例

huggingface Diffusers 将各种 Diffusion Model Pipeline 都包装好了，使用 Diffusion model 就和使用 Transformers 一样地方便：

from diffusers import DiffusionPipeline

# load model and scheduler
ldm = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")
 # run pipeline in inference (sample random noise and denoise)
prompt = "A painting of a squirrel eating a burger"

images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6).images

# save images
for idx, image in enumerate(images):
    image.save(f"squirrel-{idx}.png")

from diffusers import DiffusionPipeline

// 这里导入了 DiffusionPipeline 类，它用于方便地处理扩散模型

ldm = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")

// 使用 from_pretrained 方法加载预训练的扩散模型，指定模型名称为 "CompVis/ldm-text2im-large-256"。这一步将模型和调度器（scheduler）都准备就绪。

prompt = "A painting of a squirrel eating a burger" images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6).images

prompt 是生成图像的描述。
调用 ldm 对象，执行推理。参数说明：
num_inference_steps=50：生成图像的推理步骤数，步骤越多，图像质量通常越高。
eta=0.3：控制随机性，影响生成的图像多样性。
guidance_scale=6：引导比例，控制生成图像与提示的匹配程度，值越大，生成的图像与提示越接近。

for idx, image in enumerate(images): image.save(f"squirrel-{idx}.png")

// 遍历生成的图像列表，将每个图像保存为 PNG 格式，文件名为 squirrel-0.png，squirrel-1.png 等。

LDM Pipeline

LDM 的 pipeline 可以简化表示为：Pipeline(prompt, num_inference_steps, latents)。我们暂时考虑没有 negative prompt 和初始 latent 的输入，那么整个采样过程大致可以表示为：

1. 首先采用了 BERT 架构模型对 prompt 进行处理，生成 text_hidden_state；同时生成随机噪声 latents。

text_hidden_state = LDMBERT(prompt) # shape=[bs, len_seq, d_model] = [1, 77, 1280] 
latents = randn_tensor(latents_shape)

对于 "CompVis/ldm-text2im-large-256"，其中使用了 LDMBert， LDMBert 与传统 BERT 架构相似，规模不同，LDMBert 采用 32 层， hidden_size 为 1280，属实比 bert-base 大上不少。同时文本被 padding 到了固定的 77 长度，以此来保证文字的 hidden state 格式为 [batch_size, 77, 1280]。

2. 之后进行传统的扩散模型 backward process：

for t in self.progress_bar(self.scheduler.timesteps):
    noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample
    # compute the previous noisy sample x_t -> x_t-1
    latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample

Unet: UNet2DConditionModel：传统的 UNet 主要用于图像生成和分割, 而 UNet2DConditionModel 在此基础上增加了 Cross Attention 机制，用来综合处理文本和图像信息，使得生成的图像可以更好地符合输入的文本描述。
Scheduler：调度器（scheduler）负责控制扩散过程中的噪声添加和去噪过程。可以选择不同的调度算法，如 DDIM，以实现不同的生成效果。

for t in self.progress_bar(self.scheduler.timesteps):

// 这行代码通过一个进度条（self.progress_bar）遍历调度器（scheduler）定义的时间步（timesteps）。每个时间步 t 表示当前的扩散状态

noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample

使用 UNet 模型进行噪声预测。这里的 self.unet 是一个 UNet2DConditionModel，它不仅处理图像数据，还结合了文本信息（通过 encoder_hidden_states=context）。
latents_input 是当前的潜在样本（latent sample），t 是当前时间步。
.sample 方法返回 UNet 生成的噪声预测。

latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample

调用调度器的 step 方法，传入当前噪声预测 noise_pred、当前时间步 t 和当前潜在样本 latents。
**extra_kwargs 允许传递额外参数，以满足调度器的需求。
step 方法返回一个对象，其中 prev_sample 是计算得到的前一个噪声样本（即从 x_t 到 x_{t-1} 的转换）。

3. 最后对 latent hidden state 进行 decode，生成图片：

latents = 1 / self.vqvae.config.scaling_factor * latents
image = self.vqvae.decode(latents).sample

LDM 中的 UNET

backward process 中的 self.unet(...)，即 UNET2DCondition(sample, timestep, encoder_hidden_state) 前向推导可以看成五部分，（以下以 CompVis/ldm-text2im-large-256 为例介绍）：

准备时间步 time steps

Timesteps 编码信息是 diffusion 中 predict noise residual 模型的标配：

# 经过两次映射得到 timesteps 对应的 embedding
t_emb = self.time_proj(timesteps)
emb = self.time_embedding(t_emb, timestep_cond)

时间步编码：在扩散模型中，时间步的编码（timesteps）是预测噪声残差（noise residual）的重要信息。
self.time_proj(timesteps)：这个方法将输入的时间步 timesteps 映射到一个嵌入空间，生成 t_emb。
self.time_embedding(t_emb, timestep_cond)：进一步处理 t_emb，结合条件信息 timestep_cond，最终得到时间步的嵌入向量 emb。这个嵌入用于指导模型如何处理不同的时间步。

预处理阶段 pre-process

LDM 只用了一个 2D 卷积对输入的 hidden state 进行处理

sample = nn.Conv2d(
            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
        )(sample)

这里使用一个 2D 卷积层对输入的隐藏状态（hidden state）进行处理。
nn.Conv2d 是 PyTorch 的卷积层，用于提取特征。
- in_channels：输入通道数，指输入特征的深度。
- block_out_channels[0]：输出通道数，通常指定为模型的第一层输出。
- kernel_size 和 padding：卷积核的大小和填充方式，影响特征提取的窗口和边缘处理。

下采样过程 down sampling

down sampling 包括了 3 个 CrossAttnDownBlock2D, 和 1 个 DownBlock2D。

# down sampling 大致前向推导
down_block_res_samples = (sample,)
for downsample_block in self.down_blocks:
    sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
    # 用于 UNET 的残差链接
    down_block_res_samples += res_samples

down_block_res_samples = (sample,)

创建一个元组 down_block_res_samples，初始时包含输入样本 sample，用于存储下采样过程中产生的残差样本。

for downsample_block in self.down_blocks:
sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
down_block_res_samples += res_samples

遍历 self.down_blocks 中的每个下采样块。
调用每个下采样块，并传入当前的隐藏状态 sample、时间步嵌入 emb 和 lora_scale（可能用于调整规模）。
downsample_block 返回新的样本和残差样本，后者被追加到 down_block_res_samples 中。

其中每个 CrossAttnDownBlock2D 大概前向过程为：

# CrossAttnDownBlock2D
def forward(self, hidden_states, temb, encoder_hidden_states=None)
	output_states = ()
    for resnet, attn in zip(self.resnets, self.attentions):
        hidden_states = resnet(hidden_states, temb)
        hidden_states = attn(
            hidden_states,
            encoder_hidden_states=encoder_hidden_states,
            cross_attention_kwargs=cross_attention_kwargs,
        ).sample
        output_states += (hidden_states,)

    # downsampler = Conv2D 
    hidden_states = downsampler(hidden_states)
    output_states += (hidden_states,)

    return hidden_states, output_states

def forward(self, hidden_states, temb, encoder_hidden_states=None):
output_states = ()

output_states 用于存储每个阶段的隐藏状态。

for resnet, attn in zip(self.resnets, self.attentions):
hidden_states = resnet(hidden_states, temb)
hidden_states = attn(
hidden_states,
encoder_hidden_states=encoder_hidden_states,
cross_attention_kwargs=cross_attention_kwargs,
).sample
output_states += (hidden_states,)

对每个 ResNet 和注意力层进行迭代。
使用 ResNet 对隐藏状态 hidden_states 进行处理，结合时间步嵌入 temb。
然后通过注意力层对 hidden_states 进行处理，结合 encoder_hidden_states 进行交叉注意力（cross attention）。
将每个步骤的隐藏状态追加到 output_states 中。

hidden_states = downsampler(hidden_states)
output_states += (hidden_states,)

使用卷积层（downsampler）对隐藏状态进行下采样，将结果追加到 output_states。

return hidden_states, output_states

返回最终的 hidden_states 和 output_states。

在 CompVis/ldm-text2im-large-256 中，每个 CrossAttnDownBlock2D 包含了 2 个 attn（Transformer2DModel）以及 2 个 resnet （ResnetBlock2D）。

文字与图像的交互就发生在 Transformer2DModel 当中。每个 Transformer2DModelopen in new window 先对输入的图像数据进行预处理，将图片格式从如 (batch_size, channel, width, height) 或 (batch_size, num_image_vectors) 转换为 (batch_size, len_seq, hidden_size)，而后将 hidden_states 传入 1 层传统 Transformer layer（非 bert 或 GPT 类型），先对图像 hidden_states 进行 self-attention，而后结合 encoder_hidden_states 进行 cross attention 处理。
Transformer2DModel 中的文字与图像交互：
- 每个 CrossAttnDownBlock2D 包含两个注意力层和两个 ResNet 层。注意力层用于实现图像和文本之间的交互。
- 输入的图像数据格式会被转换，以适应 Transformer 的输入需求，通常变为 (batch_size, len_seq, hidden_size)。
- 然后，图像隐藏状态会通过自注意力（self-attention）和交叉注意力（cross attention）进行处理，结合文本信息，增强图像生成的上下文理解。

中间处理 mid processing

sample = MidBlock2DCrossAttn()(sample, 
                              emb,
                           encoder_hidden_states)

在 CompVis/ldm-text2im-large-256 中，upsampling 和 down sampling 之间采用 MidBlock2DCrossAttnopen in new window 连接，MidBlock2DCrossAttn 包括了 1 个 1 层的 Transformer2DModel 以及 1 个 resnet ResnetBlock2D。

MidBlock2DCrossAttn 是一个处理模块，位于上采样和下采样之间。
它包含一个 1 层的 Transformer2DModel 和一个 ResnetBlock2D，用于处理输入样本 sample。
emb 是时间步的嵌入向量，encoder_hidden_states 是来自文本编码器的隐藏状态，提供上下文信息。

上采样 upsampling

upsampling 采用的模块 UpBlocks 包括了 ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")，各个模块的架构与 down sampling 中的模块相似。

# upsample_block
for i, upsample_block in enumerate(self.up_blocks):
    sample = upsample_block(
                    hidden_states=sample,
                    temb=emb,
                    res_hidden_states_tuple=res_samples,
                    upsample_size=upsample_size,
                    scale=lora_scale,
                )

遍历 self.up_blocks 中的每个上采样块。
每个 upsample_block 处理当前的隐藏状态 sample，结合时间步嵌入 emb、残差隐藏状态 res_samples、上采样尺寸 upsample_size 和缩放因子 lora_scale。
模块可能包括类似于下采样的结构，如 UpBlock2D 和多个 CrossAttnUpBlock2D，这些模块结合注意力机制和卷积操作，逐步恢复图像的分辨率。

后处理 post-process

# GroupNorm
sample = self.conv_norm_out(sample)
# Silu
sample = self.conv_act(sample)
# Conv2d(320, 4, kernel=(3,3), s=(1,1), padding=(1,1))
sample = self.conv_out(sample)

总结起来，down sampling，midprocess，upsampling 三个步骤中都涉及到了 Transformer2DModel ，实现多模态的信息交互。

归一化：使用 self.conv_norm_out(sample) 进行归一化，通常是 GroupNorm，以稳定训练过程和提高生成效果。
激活函数：通过 self.conv_act(sample) 应用激活函数（如 Silu），引入非线性因素，增强模型的表达能力。
卷积输出：最后的 self.conv_out(sample) 是一个卷积层，将隐藏状态转换为最终输出图像，通常输出通道数为 4，代表生成的图像的像素通道。

LDM Super Resolution Pipeline

low_res_img = Image.open(BytesIO(response.content)).convert("RGB")
low_res_img = low_res_img.resize((128, 128)

upscaled_image = pipeline(low_res_img, num_inference_steps=100, eta=1).images[0]
upscaled_image.save("ldm_generated_image.png")

大致前项推导流程可以概括为：根据输入图片大小，生成对应的latent噪音以及time step embedding：

latents = randn_tensor(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)  # shape 与输入图片相同
latents = latents * self.scheduler.init_noise_sigma

将 latent 与原始图片拼接，然后进行 diffusion 反向推导：

for t in self.progress_bar(timesteps_tensor):
    # concat latents and low resolution image in the channel dimension.
    latents_input = torch.cat([latents, image], dim=1)
    latents_input = self.scheduler.scale_model_input(latents_input, t)
    # predict the noise residual
    noise_pred = self.unet(latents_input, t).sample
    # compute the previous noisy sample x_t -> x_t-1
    latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample

使用 vqvae 对 latent 进行解码，得到最终图片

# decode the image latents with the VQVAE
image = self.vqvae.decode(latents).sample
image = torch.clamp(image, -1.0, 1.0)
image = image / 2 + 0.5
image = image.cpu().permute(0, 2, 3, 1).numpy()

Stable diffusion

SD v1 架构

参考 hugging face diffuser 的 SD pipeline 实现open in new window。以 stable-diffusion-v1-5 为例。

1.Text Encoder

采用 CLIPTextModel，来自于 CLIPopen in new window 的 Text Encoder 部分。相比于其他传统的 Transformer 语言模型，CLIP 在预训练时，在 text-image pair 数据集上进行了对比学习预训练。prompt_embeds, negative_prompt_embeds 在经过编码后，shape 都为 [batch_size, 77, 768]

2.Diffusion 反向采样过程

SD v1.5 采样过程与 LDM 相似，其中的 latents 大小为 [bs, 4, 64, 64]。对于 txt2img，latents 通过随机生成，对于 img2img，latents 通过 VAE 模型进行 encode。

Unet 配置与 LDM 相似：

down sampling 采用 3 个 CrossAttnDownBlock2D, 和 1 个 DownBlock2D。
mid block 采用 1 个 MidBlock2DCrossAttn。hidden size = 1280
Up sampling 采用 1 个 UpBlock2D + 3 个 CrossAttnUpBlock2D

每个 CrossAttn 的 transformer 中， text embedding 大小为 768，但 Transformer 模块的 hidden size 随着 Unet 深入而增加。如 down sampling 采用的维度为 320, 640, 1280, 1280。那么 3 个 Transformer 模块中的 hidden size 就分别是 320, 640, 1280。

以 down sampling 为例，在进行 cross attention 时候，图像的 hidden state （latent）大小分别被映射到了 [4096, 320]，[2014, 640]，[256, 1280] ，而后与文字的 hidden state [77, 768] 进行 cross attention 计算。（以上张量维度省略了 batch size）

# hidden size 为 320 时候的 cross attention 单元示例
Attention(
(to_q): LoRACompatibleLinear(in_features=320, out_features=320, bias=False)
(to_k): LoRACompatibleLinear(in_features=768, out_features=320, bias=False)
(to_v): LoRACompatibleLinear(in_features=768, out_features=320, bias=False)
)

这也是 SD Unet 中 Transformer2DBlock 与传统 Transformer 主要的不同，SD Unet 中的 Transformer2DBlock 输入与输出维度是不一样的。

3.Super resolution

生成后 latent 大小为 64 * 64，通过 VQModel 解码为 512*512

SD v1.1 - v1.5

stable diffusion 1.1-1.5 的模型架构相同，以下搬运 runwaymlopen in new window 的 stable diffusion weights 总结：

sd-v1-1.ckptopen in new window: 237k steps at resolution 256x256 on laion2B-enopen in new window. 194k steps at resolution 512x512 on laion-high-resolutionopen in new window (170M examples from LAION-5B with resolution >= 1024x1024).
sd-v1-2.ckptopen in new window: Resumed from sd-v1-1.ckpt. 515k steps at resolution 512x512 on laion-aesthetics v2 5+open in new window (a subset of laion2B-en with estimated aesthetics score > 5.0, and additionally filtered to images with an original size >= 512x512, and an estimated watermark probability < 0.5. The watermark estimate is from the LAION-5Bopen in new window metadata, the aesthetics score is estimated using the LAION-Aesthetics Predictor V2open in new window).
sd-v1-3.ckptopen in new window: Resumed from sd-v1-2.ckpt. 195k steps at resolution 512x512 on "laion-aesthetics v2 5+" and 10% dropping of the text-conditioning to improve classifier-free guidance samplingopen in new window.
sd-v1-4.ckptopen in new window: Resumed from sd-v1-2.ckpt. 225k steps at resolution 512x512 on "laion-aesthetics v2 5+" and 10% dropping of the text-conditioning to improve classifier-free guidance samplingopen in new window.
sd-v1-5.ckptopen in new window: Resumed from sd-v1-2.ckpt. 595k steps at resolution 512x512 on "laion-aesthetics v2 5+" and 10% dropping of the text-conditioning to improve classifier-free guidance samplingopen in new window.
sd-v1-5-inpainting.ckptopen in new window: Resumed from sd-v1-5.ckpt. 440k steps of inpainting training at resolution 512x512 on "laion-aesthetics v2 5+" and 10% dropping of the text-conditioning to improve classifier-free guidance samplingopen in new window. For inpainting, the UNet has 5 additional input channels (4 for the encoded masked-image and 1 for the mask itself) whose weights were zero-initialized after restoring the non-inpainting checkpoint. During training, we generate synthetic masks and in 25% mask everything.

SD v2

参考 stability-AI 仓库open in new window，SD v2 相对 v1 系列改动较大：

架构方面 SD v2 系列：

采用了 OpenCLIP-ViT/Hopen in new window 作为 text encoder。
Unet 架构改变：其中 Transformer 模块中的 attention_head_dim 变为了 5,10,20,20，SD v1 中为 8,8,8,8。cross_attention_dim 从 768 变为 1280。同时在 latent hidden state 进入 cross attention 之前，额外采用了 linear_projection 进行 latent hidden state 的处理，SD v1 中为卷积层处理。

训练方面 SD v2 系列，（以下拷贝了 huggingface 中 SD 模型 model card 的介绍）：

SD 2.0-baseopen in new window：The model is trained from scratch 550k steps at resolution 256x256 on a subset of LAION-5Bopen in new window filtered for explicit pornographic material, using the LAION-NSFW classifieropen in new window with punsafe=0.1 and an aesthetic scoreopen in new window >= 4.5. Then it is further trained for 850k steps at resolution 512x512 on the same dataset on images with resolution >= 512x512.
SD v2.0open in new window：This stable-diffusion-2 model is resumed from stable-diffusion-2-baseopen in new window (512-base-ema.ckpt) and trained for 150k steps using a v-objectiveopen in new window on the same dataset. Resumed for another 140k steps on 768x768 images.
SD v2.1open in new window：This stable-diffusion-2-1 model is fine-tuned from stable-diffusion-2open in new window (768-v-ema.ckpt) with an additional 55k steps on the same dataset (with punsafe=0.1), and then fine-tuned for another 155k extra steps with punsafe=0.98.

Lora

huggingface diffuser 中 Lora 的实现与 huggingface/PEFT 实现方法相似，添加 Lora 只需要通过撰写规则，锁定需要改动的 layer，并替换为 LoRACompatibleLayer 实现，huggingface 也提供好了 lora 训练代码open in new window，和 SD lora 推理方法。

Diffusers 中，SD 采用 Lora 的部分位于 Unet 当中，大部分的 Lora 在 Transformer 模块当中，SD 的 lora 与 NLP Lora 实现方式基本相同， 一个较大的区别在于，SD 中的 Lora 除了对线性层进行 Lora 叠加外，也对卷积层进行了 Lora 改造 。

绝大部分内容搬运自Kevin 吴嘉文，自己补充了一些细节方面

DIFFUSION 系列笔记| Latent Diffusion Modelhttps://antarina.tech/posts/notes/articles/%E7%AC%94%E8%AE%B0latent_diffusion.html#sd-v2