目录
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-1latents = 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,).sampleoutput_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 residualnoise_pred = self.unet(latents_input, t).sample# compute the previous noisy sample x_t -> x_t-1latents = 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 resolution512x512
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 resolution512x512
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 resolution512x512
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 resolution512x512
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 resolution512x512
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 resolution512x512
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 withpunsafe=0.1
and an aesthetic scoreopen in new window >=4.5
. Then it is further trained for 850k steps at resolution512x512
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 on768x768
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 (withpunsafe=0.1
), and then fine-tuned for another 155k extra steps withpunsafe=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