系列文章目录
- 【扩散模型(一)】中介绍了 Stable Diffusion 可以被理解为重建分支(reconstruction branch)和条件分支(condition branch)
- 【扩散模型(二)】IP-Adapter 从条件分支的视角,快速理解相关的可控生成研究
- 【可控图像生成系列论文(一)】 简要介绍了 MimicBrush 的整体流程和方法;
- 【可控图像生成系列论文(二)】 就MimicBrush 的具体模型结构、训练数据和纹理迁移进行了更详细的介绍。
- 【可控图像生成系列论文(三)】介绍了一篇相对早期(2018年)的可控字体艺术化工作。
- 【可控图像生成系列论文(四)】介绍了 IP-Adapter 具体是如何训练的?
- 【可控图像生成系列论文(五)】ControlNet 和 IP-Adapter 之间的区别有哪些?
- 本文《【扩散模型(三)】IP-Adapter 源码详解1-输入篇》作为两个系列的交汇点,将通过对经典的 IP-Adapter 源码详细阅读,进一步加深对其原理的解释。
文章目录
- 系列文章目录
- 整体结构图+代码中的变量名
- 一、IP-Adapter 做了什么?
- 二、对应的代码实现
- 1.模型输入
- 2.Linear 和 LN(LayerNorm)
- 总结
整体结构图+代码中的变量名
IP-Adapter 源码:https://github.com/tencent-ailab/IP-Adapter
本文就基于 SD1.5 的 IP-Adapter 训练代码 tutorial_train.py 为例,进行代码和结构图的解释。
一、IP-Adapter 做了什么?
如上图所示,插入了图中的最上面一条分支(图像输入条件分支):
- 蓝色的(无需训练的) Image Encoder
- 红色的(需训练的)Linear + LN(LayerNorm)
- 红色的(需训练的)、针对图像(Image Prompt)的 Cross Attention。
在论文中也提到,具体分别是:
- Image Encoder 是 pretrained CLIP image encoder
- 线性层和层归一化 Linear + LN(LayerNorm1):
- 为了有效地分解全局图像嵌入,作者使用一个小的可训练投影网络(projection network)将图像嵌入投影到长度为N的特征序列中(在本研究中使用N=4),图像特征的维数与预训练的扩散模型中文本特征的维数相同。使用的投影网络由线性层和层归一化组成。
- Decoupled Cross-Attention 中,做法是在原来的 UNet 的 Cross-Attention 中加了一层 Cross-Attention。
- 如原文提到 “we add a new cross-attention layer for each cross-attention layer in the original UNet model to insert image features.”
二、对应的代码实现
1.模型输入
先简单看下模型的训练时的输入,即 /path/IP-Adapter/tutorial_train.py 中 main() 函数内的 dataloader 部分,下面代码通过调用 MyDataset 类来实现了 train_dataloader 的构建。
# dataloader
train_dataset = MyDataset(args.data_json_file, tokenizer=tokenizer, size=args.resolution, image_root_path=args.data_root_path)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
shuffle=True,
collate_fn=collate_fn,
batch_size=args.train_batch_size,
num_workers=args.dataloader_num_workers,
)
对于实际训练使用的数据则为从 train_dataloader 中取的:
- batch[“images”]
- 用来得到形状后,生成随机噪声。
- 具体如下代码所示,通过
vae.encoder
得到 latents后 - 通过
torch.randn_like(latents)
按照latents
张量的形状生成一个随机的噪声张量noise
。
- batch[“clip_images”]
- 通过
image_encoder
得到image_embeds
图像特征
- 通过
- batch[“drop_image_embeds”]
- 文中有提到会随机通过随机丢弃条件信息(如文本或图像嵌入),使得模型会学会在有条件和无条件的情况下进行预测(生成图像)
- batch[“text_input_ids”] 是文本输入,通过一个 text_encoder 后得到文本特征
encoder_hidden_states
。
for step, batch in enumerate(train_dataloader):
load_data_time = time.perf_counter() - begin
with torch.no_grad():
latents = vae.encode(batch["images"].to(accelerator.device, dtype=weight_dtype)).latent_dist.sample()
latents = latents * vae.config.scaling_factor
# Sample noise that we'll add to the latents
noise = torch.randn_like(latents)
bsz = latents.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device)
timesteps = timesteps.long()
# Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
with torch.no_grad():
image_embeds = image_encoder(batch["clip_images"].to(accelerator.device, dtype=weight_dtype)).image_embeds
image_embeds_ = []
for image_embed, drop_image_embed in zip(image_embeds, batch["drop_image_embeds"]):
if drop_image_embed == 1:
image_embeds_.append(torch.zeros_like(image_embed))
else:
image_embeds_.append(image_embed)
image_embeds = torch.stack(image_embeds_)
with torch.no_grad():
encoder_hidden_states = text_encoder(batch["text_input_ids"].to(accelerator.device))[0] # pooled_prompt_embeds?
2.Linear 和 LN(LayerNorm)
以 SD1.5 + IP-Adapter 的训练代码为例:
下方代码为 /path/IP-Adapter/tutorial_train.py 中 main() 函数内,调用了定义好的 ImageProjModel 类
#ip-adapter
image_proj_model = ImageProjModel(
cross_attention_dim=unet.config.cross_attention_dim,
clip_embeddings_dim=image_encoder.config.projection_dim,
clip_extra_context_tokens=4,
)
下方代码为 /path/IP-Adapter/ip_adapter/ip_adapter.py 被调用的 ImageProjModel 类,在构造函数 __init__
中可以看到有前文提到的 Linear 和 LayerNorm。
class ImageProjModel(torch.nn.Module):
"""Projection Model"""
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
super().__init__()
self.generator = None
self.cross_attention_dim = cross_attention_dim
self.clip_extra_context_tokens = clip_extra_context_tokens
self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
self.norm = torch.nn.LayerNorm(cross_attention_dim)
def forward(self, image_embeds):
embeds = image_embeds
clip_extra_context_tokens = self.proj(embeds).reshape(
-1, self.clip_extra_context_tokens, self.cross_attention_dim
)
clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
return clip_extra_context_tokens
总结
本文详解了IP-Adapter 训练源码中的输入部分,下篇则详解核心部分,针对图像输入的 Cross-Attention。
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016 ↩︎