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

目录

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

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

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

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)

预处理阶段 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)

下采样过程 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,)

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

其中每个 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 = ()

    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,)

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

    return 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 处理。

中间处理 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。

上采样 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,
                )

后处理 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 ，实现多模态的信息交互。

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 系列：

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

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

总结