学习前言
Inpaint是Stable Diffusion中的常用方法，一起简单学习一下。
源码下载地址
https://github.com/bubbliiiing/stable-diffusion
喜欢的可以点个star噢。
原理解析
一、先验知识
txt2img的原理如博文
Diffusion扩散模型学习2——Stable Diffusion结构解析-以文本生成图像（文生图，txt2img）为例
img2img的原理如博文
Diffusion扩散模型学习3——Stable Diffusion结构解析-以图像生成图像（图生图，img2img）为例
二、什么是inpaint
Inpaint是一项图片修复技术，可以从图片上去除不必要的物体，让您轻松摆脱照片上的水印、划痕、污渍、标志等瑕疵。
一般来讲，图片的inpaint过程可以理解为两步：
1、找到图片中的需要重绘的部分，比如上述提到的水印、划痕、污渍、标志等。
2、去掉水印、划痕、污渍、标志等，自动填充图片应该有的内容。
三、Stable Diffusion中的inpaint
Stable Diffusion中的inpaint的实现方式有两种：
1、开源的inpaint模型
参考链接：inpaint_st.py，该模型经过特定的训练。需要输入符合需求的图片才可以进行inpaint。
需要注意的是，该模型使用的config文件发生了改变，改为v1-inpainting-inference.yaml。其中最显著的区别就是unet_config的in_channels从4变成了9。相比于原来的4，我们增加了4+1（5）个通道的信息。
4+1（5）个通道的信息应该是什么呢？一个是被mask后的图像，对应其中的4；一个是mask的图像，对应其中的1。
- 1、我们首先把图片中需要inpaint的部分给置为0，获得被mask后的图像，然后利用VAE编码，VAE输出通道为4，假设被mask的图像是[512, 512, 3]，此时我们获得了一个[4, 64, 64]的隐含层特征，对应其中的4。
- 2、然后需要对mask进行下采样，采样到和隐含层特征一样的高宽，即mask的shape为[1, 512, 512]，利用下采样获得[1, 64, 64]的mask。本质上，我们获得了隐含层的mask。
- 3、然后我们将 下采样后的被mask的图像 和 隐含层的mask 在通道上做一个堆叠，获得一个[5, 64, 64]的特征，然后将此特征与随机初始化的高斯噪声堆叠，则获得了上述图片中的9通道特征。
此后采样的过程与常规采样方式一样，全部采样完成后，使用VAE解码，获得inpaint后的图像。
可以感受到上述的方式必须基于一个已经训练好的unet模型，这要求训练者需要有足够的算力去完成这一个工作，对大众开发者而言并不友好。因此该方法很少在实际中得到使用。
2、基于base模型inpaint
如果我们必须训练一个inpaint模型才能对当前的模型进行inpaint，那就太麻烦了，有没有什么方法可以不需要训练就能inpaint呢？
诶诶，当然有哈。
Stable Diffusion就是一个生成模型，如果我们可以做到让Stable Diffusion只生成指定区域，并且在生成指定区域的时候参考其它区域，那么它自身便是一个天然的inpaint模型。
如何做到这一点呢？我们需要结合img2img方法，我们首先考虑inpaint的两个输入：一个是原图，另外一个是mask图。
在img2img中，存在一个denoise参数，假设我们设置denoise数值为0.8，总步数为20步，那么我们会对输入图片进行0.8x20次的加噪声。如果我们可以在这个加噪声图片的基础上进行重建，那么网络必然会考虑加噪声图（也就对应了原始图片的特征）。
在图像重建的20步中，对隐含层特征，我们利用mask将不重建的地方都替换成 原图按照当前步数加噪后的隐含层特征。此时不重建的地方的特征都由输入图片决定。然后不替换需要重建的地方进行，利用unet计算噪声进行重建。
具体部分，可看下面的循环与代码，我已经标注出了 替换特征的地方，在这里mask等于1的地方保留原图，mask等于0的地方不断的重建。
- 将原图x0映射到VAE隐空间，得到img_orig；
- 初始化随机噪声img（也可以使用img_orig完全加噪后的噪声）；
- 开始循环：
- 对于每一次时间步，根据时间步生成img_orig对应的噪声特征；
- 一个是基于上个时间步降噪后得到的img，一个是基于原图得到的img_orig。通过mask将两者融合， i m g = i m g _ o r i g ∗ m a s k + ( 1.0 − m a s k ) ∗ i m g img = img\_orig * mask + (1.0 - mask) * img img=img_orig∗mask+(1.0−mask)∗img。即，将原图中的非mask区域和噪声图中的mask区域进行融合，得到新的噪声图。
- 然后继续去噪声直到结束。
由于该方法不需要训练新模型，并且重建效果也不错，所以该方法比较通用。
- for i, step in enumerate(iterator):
- # index是用来取得对应的调节参数的
- index = total_steps - i - 1
- # 将步数拓展到bs维度
- ts = torch.full((b,), step, device=device, dtype=torch.long)
- # --------------------------------------------------------------------------------- #
- # 替换特征的地方
- # 用于进行局部的重建，对部分区域的隐向量进行mask。
- # 对传入unet前的隐含层特征，我们利用mask将不重建的地方都替换成 原图加噪后的隐含层特征
- # self.model.q_sample用于对输入图片进行ts步数的加噪
- # --------------------------------------------------------------------------------- #
- if mask is not None:
- assert x0 is not None
- img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
- img = img_orig * mask + (1. - mask) * img
- # 进行采样
- outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
- quantize_denoised=quantize_denoised, temperature=temperature,
- noise_dropout=noise_dropout, score_corrector=score_corrector,
- corrector_kwargs=corrector_kwargs,
- unconditional_guidance_scale=unconditional_guidance_scale,
- unconditional_conditioning=unconditional_conditioning)
- img, pred_x0 = outs
- # 回调函数
- if callback: callback(i)
- if img_callback: img_callback(pred_x0, i)
- if index % log_every_t == 0 or index == total_steps - 1:
- intermediates['x_inter'].append(img)
- intermediates['pred_x0'].append(pred_x0)
根据通用性，本文主要以上述提到的基于base模型inpaint进行解析。
1、输入图片到隐空间的编码
inpaint技术衍生于图生图技术，所以同样需要指定一张参考的图像，然后在这个参考图像上开始工作。
利用VAE编码器对这张参考图像进行编码，使其进入隐空间，只有进入了隐空间，网络才知道这个图像是什么。
此时我们便获得在隐空间的图像，后续会在这个 隐空间加噪后的图像 的基础上进行采样。
2、文本编码
文本编码的思路比较简单，直接使用CLIP的文本编码器进行编码就可以了，在代码中定义了一个FrozenCLIPEmbedder类别，使用了transformers库的CLIPTokenizer和CLIPTextModel。
在前传过程中，我们对输入进来的文本首先利用CLIPTokenizer进行编码，然后使用CLIPTextModel进行特征提取，通过FrozenCLIPEmbedder，我们可以获得一个[batch_size, 77, 768]的特征向量。
- class FrozenCLIPEmbedder(AbstractEncoder):
- """Uses the CLIP transformer encoder for text (from huggingface)"""
- LAYERS = [
- "last",
- "pooled",
- "hidden"
- ]
- def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
- freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32
- super().__init__()
- assert layer in self.LAYERS
- # 定义文本的tokenizer和transformer
- self.tokenizer = CLIPTokenizer.from_pretrained(version)
- self.transformer = CLIPTextModel.from_pretrained(version)
- self.device = device
- self.max_length = max_length
- # 冻结模型参数
- if freeze:
- self.freeze()
- self.layer = layer
- self.layer_idx = layer_idx
- if layer == "hidden":
- assert layer_idx is not None
- assert 0 <= abs(layer_idx) <= 12
- def freeze(self):
- self.transformer = self.transformer.eval()
- # self.train = disabled_train
- for param in self.parameters():
- param.requires_grad = False
- def forward(self, text):
- # 对输入的图片进行分词并编码，padding直接padding到77的长度。
- batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
- return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
- # 拿出input_ids然后传入transformer进行特征提取。
- tokens = batch_encoding["input_ids"].to(self.device)
- outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer=="hidden")
- # 取出所有的token
- if self.layer == "last":
- z = outputs.last_hidden_state
- elif self.layer == "pooled":
- z = outputs.pooler_output[:, None, :]
- else:
- z = outputs.hidden_states[self.layer_idx]
- return z
- def encode(self, text):
- return self(text)
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！ |
