HD-Painter: 基于扩散模型的高分辨率实时文本引导图像修复

用户1324186

发布于 2024-01-04 13:04:48

3260

发布于 2024-01-04 13:04:48

来源：arxiv 作者：Hayk Manukyan,Andranik Sargsyan 等 论文题目：HD-Painter: High-Resolution and Prompt-Faithful Text-Guided Image Inpainting with Diffusion Models 论文链接：https://arxiv.org/pdf/2312.14091.pdf 项目主页：https://github.com/Picsart-AI-Research/HD-Painter 内容整理：汪奕文当前文本-图像 inpainting 模型仍有很大的改进潜力，尤其是在更好地将 inpainting 区域与用户图像对齐以及执行高分辨率 inpainting 方面。因此，本文中提出了 HD-Painter，一种无需训练的方法，可以准确地根据提示并连贯地扩展到高分辨率图像 inpainting 。

引言

经过预训练的文本到图像生成模型，如 SD、Imagen 和 Dall-E 2，可以在后向扩散过程中将扩散的已知区域与生成（去噪）的未知区域混合，从而完成图像补全。虽然这些方法能生成和谐且视觉上合理的补全图像，但它们缺乏对全局场景的理解，提示忽略的主要缺点表现在两种情况下：

背景主导： 当图像的未知区域主要以背景元素填充时，会忽略文本提示的具体内容。
邻近物体主导： 已知区域的物体根据视觉上下文可能性而不是给定的文本提示传播到未知区域。

出现这两个问题的原因可能是，扩散 inpainting 缺乏准确解释文本提示或将其与已知区域的上下文信息相结合的能力。为了解决上述问题，我们引入了无需任何训练或微调的提示感知内向注意（PAIntA）模块，可根据给定的文本条件增强自我注意力得分，旨在减少图像已知区域中与提示无关信息的影响，同时增加与提示对齐的已知像素的贡献。

本文的主要贡献如下：

引入了 "提示感知内向注意力"（PAIntA）层，以缓解文本引导的图像 inpainting 中背景和附近物体占主导地位的提示忽略问题。
为了进一步改进文本对齐生成，我们提出了重新权重注意力分数指导（RASG）策略，该策略能够在指导采样时，防止出现分布偏移。
本文提出的文本引导的图像补全方法完全无需训练，与目前最先进的方法相比，在定量和定性方面都具有显著优势。此外，在简单而有效的 inpainting 专用超分辨率框架的额外帮助下，我们还能完成高分辨率（高达

2048\times2048

）的图像 inpainting。

方法

I\in \mathbb{R}^{H\times W\times 3}

是 RGB 图像，

M\in\left\{0, 1\right\}^{H \times W}

是一个二进制 mask，表示在

中要用文本提示

\tau

进行 inpainting 的区域。文本引导图像绘制的目标是输出图像

I^c \in \mathbb{R}^{H\times W\times 3}

，使

I^c \in \mathbb{R}^{H\times W\times 3}

包含 mask 区域内提示

\tau

所描述的对象，而 mask 区域外则与

重合，即

I^c\odot (1-M) = I \odot (1-M)

。

图1

本文提出的 pipeline 由两个阶段组成：在

H/4 \times W/4

分辨率上应用文本引导的图像 inpainting，然后对生成的内容进行

\times 4

超分辨率。

为了根据给定的文本提示

\tau

完成缺失区域

的 inpainting，我们采用了类似 SD inpainting 的预训练扩散模型，用 PAIntA 层替换自注意层，并通过应用 RASG 机制执行扩散后向处理。在得到最终估计的隐空间特征

x_0

后，对其进行解码，得到 inpainted 图像

I^{low}_c = \mathcal{D}(x_0) ∈ \mathbb{R}^{\frac{H}{4} \times \frac{W}{4}}

。

为了对原始大小的图像

I\in \mathbb{R}^{H\times W}

进行 inpainting，使用了 LDM，从

X_T\sim \mathcal{N}(0, 1)

开始应用 SD 的后向扩散过程，并以低分辨率 inpainted 图像

I^{low}_c

为条件。每一步之后，将去噪后的

X^{pred}_0

与 mask

(1-M)\in \left\{0, 1\right\}^{H\times W}

所指示的已知区域中的原始图像编码

\mathcal{E}(I)

融合，得到下一个隐空间特征

X_{t-1}

。在最后一步之后，通过

\mathcal{D}(X_0)

对隐空间特征解码，并使用泊松混合法对

进行融合，以避免边缘伪影。

提示感知内向注意

令

X\in \mathbb{R}^{(h\times w)\times d}

为 PAIntA 的输入。与自注意力类似，PAIntA 首先得到

Q_s,K_s,V_s \in \mathbb{R}^{(h \times w)\times d}

以及相似度矩阵

A_{self} = \frac{Q_sK_s^T}{\sqrt{d}}\in \mathbb{R}^{hw\times hw}

。

然后，通过调整对 inpainting 区域有贡献的已知像素的注意力分数，来减轻已知区域对未知区域过强的影响。具体来说，利用文本提示

\tau

，PAIntA 定义了一个新的相似度矩阵：

(\tilde{A}_{self})_{ij}=\left\{ \begin{aligned} &c_j \cdot (A_{self})_{ij} &M_j=0 ~\text{and}~M_j=1 ,\\ &(A_{self})_{ij} &\text{otherwise.} \end{aligned} \right. \quad(1)

其中，

c_j

表示第

个特征 token（像素）与给定文本提示

\tau

的对齐度。

用交叉注意力空间文本相似性矩阵

S_{cross} = \text{SoftMax}(Q_cK_c^T /\sqrt{d})

来定义

\left\{c_j\right\}^{hw}_{j=1}

，其中

Q_c\in \mathbb{R}^{(h \times w)\times d},K_c\in \mathbb{R}^{l\times d}

。具体来说，考虑提示文本

\tau

的 CLIP 文本嵌入，并分离出与

\tau

的单词和 EOT token 相对应的嵌入，并用

ind(\tau) \subset \left\{1,2,\dots,l\right\}

表示所选索引集。对于第

个像素，将其与文本提示

\tau

的相似度分数相加，

c_j = \sum_{k\in ind(\tau)}(S_{cross})_{jk}

。

图2

需要注意的是，在 vanilla SD 中，交叉注意层位于自注意层之后，因此在 PAIntA 中，要获得

Q_c,K_c

，需要借用下一个交叉注意模块的投影层权重。

重新权重注意力分数指导

为了进一步提高生成与提示

\tau

的一致性，采用了一种采样后指导的机制，其目标函数

S(x)

利用了交叉注意层的开放词汇分割特性。具体来说，每一步中，在预测噪声

\epsilon_{\theta}^t(x_t)

后使用以下更新规则：

\hat{\epsilon}^t_{\theta}(x_t) \leftarrow \epsilon^t_{\theta}(x_t)+\sqrt{1 - \alpha_t} \cdot s \nabla_{x_t} S(x_t)

，其中

是控制引导量的超参数。将DDIM 过程中的

\epsilon^t_{\theta}(x_t)

替换为

\hat{\epsilon}^t_{\theta}(x_t)

后，可以得到

x_{t-1} = \sqrt{\alpha_{t-1}}\frac{x_t - \sqrt{1 - \alpha_t} \epsilon^t_{\theta}(x_t)}{\sqrt{\alpha_{t}}} + {\sqrt{1 - \alpha_{t-1} }\epsilon^t_{\theta}(x_t) - \xi_t \nabla_{x_t} S(x_t)},\\ \xi_t = \sqrt{1 - \alpha_t} \cdot s \left( \frac{\sqrt{1 - \alpha_t} \sqrt{\alpha_{t-1}}}{\sqrt{\alpha_t}} - {\sqrt{1 - \alpha_{t-1}}}\right)\quad(2)

因此，在公式 2 中，得到了一个附加项

- \xi_t \nabla_{x_t} S(x_t)

，它可能会改变

x_{t-1}

的原始分布。为此，我们引入了重新加权注意力分数引导（RASG）策略，引入了梯度重新加权机制，从而实现隐空间保护。具体来说，在公式中用梯度

\nabla_{x_t} S(x_t)

的重定向版本（使其更接近于从

\mathcal{N} (0, 1)

中采样）来代替随机分量

\epsilon_t

，就能使

x_{t-1}

保持在所需的域内，同时引导采样过程最小化

S(x_t)

。

梯度

\nabla_{x_t} S(x_t)

的缩放是通过除以标准差来实现的。因此，RASG 采样的公式为

x_{t-1} = \sqrt{\alpha_{t-1}}\frac{x_t - \sqrt{1 - \alpha_t} \epsilon_{\theta}(x_t)}{\sqrt{\alpha_{t}}} +\sqrt{1 - \alpha_{t-1}-\sigma_t^2} \epsilon_{\theta}(x_t) +\sigma_t\frac{\nabla_{x_t} S(x_t)}{\text{std}(\nabla_{x_t} S(x_t))}\quad(3)