【扩散模型（三）】IP-Adapter 源码详解1-输入篇

系列文章目录

【扩散模型（一）】中介绍了 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（LayerNorm¹）：
- 为了有效地分解全局图像嵌入，作者使用一个小的可训练投影网络（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 的构建。

    # dataloadertrain_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() - beginwith 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 latentsnoise = torch.randn_like(latents)bsz = latents.shape[0]# Sample a random timestep for each imagetimesteps = 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_embedsimage_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-adapterimage_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 = Noneself.cross_attention_dim = cross_attention_dimself.clip_extra_context_tokens = clip_extra_context_tokensself.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_embedsclip_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