4、解构三个重要的Pipeline(SD-Inpainting, ControlNet, AnimateDiff) [代码级手把手解析diffusers库]_animatediff diffuser

上一篇我们解析了所有Pipeline的基类 DiffusionPipeline。后续各种各样的pipeline都继承了DiffusionPipeline的模型加载保存等功能，然后再配合各个组件实现各种的结构即可。

事实上，一个Pipeline通常包含了如下模块(from_pretrained函数根据model_index.json文件new了一个Pipeline，挨个扫描子文件夹创建模块加载对应的weight和config)：

• VAE，即变分自编码器，把图像编码到特征，进行生成过程后再把特征解码到图像。(有weight)
• UNet，用于迭代采样预测噪声的模型。(有weight)
• Text Encoder，用于把tokens编码为一串向量，用来控制扩散模型的生成。(有weight)
• Tokenizer，把输入的文本按照字典编码为上面的tokens。
• Scheduler，我们知道扩散模型有很多采样方法，Scheduler定义了我们用哪种采样方法
• Safety_checker，NSFW检测器，很多人应该都不想要这个，可以去掉。
• Feature_extractor，也是NSFW检测器的一部分，也可以去掉。
• Image Encoder，如果有image作为条件就会需要，计算image embedding，用来控制扩散模型的生成(有weight)
• Image Processor，配合Image Encoder，在计算emebedding之前进行一些数据增强。
• ControlNet，如果使用各种control condition就会需要，本质是半个UNet，用于计算condition_images的feature 融合到UNet中(有weight)

本节我将带领大家解构diffusers中3个我觉得最有学习意义的Pipeline：

• SD-Inpainting：StableDiffusionInpaintPipeline
• ControlNet：StableDiffusionControlNetPipeline
• AnimateDiff：AnimateDiffPipeline

理解了上方这几个Pipeline，我们后续就很好对任意的pipeline进行修改和自定义了，无论是需要图像生成结构控制，还是局部重绘，亦或是需要视频生成一致性控制，都可以有一个很好的理解。

对于每个pipeline，我将按照如下方式展开：

• 首先，进行组件的介绍__init__。
• 然后，按照pipeline的执行顺序__call__，分模块的讲解实现代码和原理。

【注意：SD作为经典的Pipeline我放在前面讲解，后面的Pipeline的某些代码，如果和SD相同，我将不再重复讲解，如果有细微差别，我会专门讲解】

题外话：我们看那典中典的五行代码，最后一行pipeline(xxx)，代表着我们用这个对象的名称作为一个函数的调用，那么这种用法就会自动调用StableDiffusionPipeline类中的__call__()函数。事实上，Diffusers库中的大多数类我们首先就要看__init__()函数和__call__()函数。

StableDiffusionInpaintPipeline

Inpaint是一项图片局部重绘技术，可以从图片上去除不必要的物体，让您轻松摆脱照片上的水印、划痕、污渍、标志等瑕疵。

一般来讲，图片的inpaint过程可以理解为两步：
1、找到图片中的需要重绘的部分，比如上述提到的水印、划痕、污渍、标志等。
2、自动填充图片重绘区域应该有的内容，如去掉水印、划痕、污渍、标志等。

组件：

vae: Union[AutoencoderKL, AsymmetricAutoencoderKL],
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
1
2
3
4
5
6
7
8

加载模型

根据上节可知，StableDiffusionInpaintPipeline调用from_pretrained依次加载各个模块，再最后执行__init__()函数：这一段把所有模块组合起来，并为pipeline注册对应的配置信息(一个FrozenDict类self._internal_dict中)。

	def __init__(
        self,
        vae: Union[AutoencoderKL, AsymmetricAutoencoderKL],
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        unet: UNet2DConditionModel,
        scheduler: KarrasDiffusionSchedulers,
        safety_checker: StableDiffusionSafetyChecker,
        feature_extractor: CLIPImageProcessor,
        image_encoder: CLIPVisionModelWithProjection = None,
        requires_safety_checker: bool = True,
    ):
        super().__init__()

		# 中间是对safety_checker的警告，可以直接无视，
		
        self.register_modules(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            scheduler=scheduler,
            safety_checker=safety_checker,
            feature_extractor=feature_extractor,
            image_encoder=image_encoder,
        )

		# vae的缩放系数，以及vae的图像预处理操作
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
        self.mask_processor = VaeImageProcessor(
            vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
        )
        self.register_to_config(requires_safety_checker=requires_safety_checker)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

这里的所有模块已经new好并已经读取了对应的模型文件（由DiffusionPipeline.from_pretrained实现）

register_to_config完成模型参数的配置（维护一个字典_internal_dict），我们可以打印pipeline._internal_dict查看配置信息：

	FrozenDict([('vae', ('diffusers', 'AutoencoderKL')),
            ('text_encoder', ('transformers', 'CLIPTextModel')),
            ('tokenizer', ('transformers', 'CLIPTokenizer')),
            ('unet', ('diffusers', 'UNet2DConditionModel')),
            ('scheduler', ('diffusers', 'PNDMScheduler')),
            ('safety_checker',
             ('stable_diffusion', 'StableDiffusionSafetyChecker')),
            ('feature_extractor', ('transformers', 'CLIPImageProcessor')),
            ('requires_safety_checker', True)])
1
2
3
4
5
6
7
8
9

其中对于difusers库的组件(vae、unet、scheduler)，我们也可以查看它们的_internal_dict，内部保存了模型的结构和配置参数，如vae._internal_dict：

FrozenDict([('in_channels', 3),
            ('out_channels', 3),
            ('down_block_types',
             ['DownEncoderBlock2D',
              'DownEncoderBlock2D',
              'DownEncoderBlock2D',
              'DownEncoderBlock2D']),
            ('up_block_types',
             ['UpDecoderBlock2D',
              'UpDecoderBlock2D',
              'UpDecoderBlock2D',
              'UpDecoderBlock2D']),
            ('block_out_channels', [128, 256, 512, 512]),
            ('layers_per_block', 2),
            ('act_fn', 'silu'),
            ('latent_channels', 4),
            ('norm_num_groups', 32),
            ('sample_size', 512),
            ('_class_name', 'AutoencoderKL'),
            ('_diffusers_version', '0.6.0'),
            ('_name_or_path',
             '/data1/huggingface/StableDiffusion/stable-diffusion-v1-5/vae')])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

此Pipe继承自 DiffusionPipeline。除此之外Pipe还继承了以下加载方法：

• load_textual_inversion（） 用于加载文本反转嵌入
• load_lora_weights（） 用于加载 LoRA 权重
• save_lora_weights（） 用于节省 LoRA 权重
• load_ip_adapter（） 用于加载 IP 适配器
• from_single_file（） 用于加载文件.ckpt

Pipe推理流程

如何做到这一点呢？我们需要结合img2img方法，我们首先考虑inpaint的两个输入：一个是原图，另外一个是mask图。

Stable Diffusion中的inpaint的实现方式有两种，由num_channels_unet = self.unet.config.in_channels参数决定：

• 重新训练in_channel=9的UNet：只需要更改unet的输入channel=9=(4+4+1)，重新训练，直接将（4通道的noise_latent 、 4通道的masked_image_latent 、1通道的mask_latent ）三者进行concat送入unet即可。
• 利用任意预训练好的in_channel=4的UNet：让Stable Diffusion只生成指定区域，并且在生成指定区域的时候参考其它区域。在图像重建的20步中，对于每个timestep送入unet的latent特征$newLatent=mask\ast imageLatent+(1-mask)\ast oldLatent$，我们利用mask将不重建的地方都替换成 原图按照当前步数加噪后的latent特征(此时不重建的地方的特征都由输入图片决定)，然后不替换需要重建区域的latent，利用unet计算噪声更新需要重建的地方。
  
  

Pipeline的__call__函数就是整个生成推理环节的代码：

	@torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        image: PipelineImageInput = None,
        mask_image: PipelineImageInput = None,
        masked_image_latents: torch.FloatTensor = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        padding_mask_crop: Optional[
1
2
3
4
5
6
7
8
9