4、解构三个重要的Pipeline(SD-Inpainting, ControlNet, AnimateDiff) [代码级手把手解析diffusers库]
上一篇我们解析了所有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
组件:
vae: Union[AutoencoderKL, AsymmetricAutoencoderKL],
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
加载模型
根据上节可知,StableDiffusionInpaintPipeline调用from_pretrained
依次加载各个模块,再最后执行__init__()
函数:这一段把所有模块组合起来,并为pipeline注册对应的配置信息(一个FrozenDict
类self._internal_dict
中)。
def __init__(
self,
vae: Union[AutoencoderKL, AsymmetricAutoencoderKL