跳转至

UNet-Twice: A Simple Structured Reference-based Inpainting Framework

会议: ECCV 2026
论文: ECCV 原文
领域: 图像生成
关键词: 参考图像修复 / 扩散模型 / 权重复用 / 互补掩码 / 图像修复

一句话总结

提出 UNet-Twice 框架,通过在单次前向过程中复用同一个扩散 UNet 两次并结合互补掩码与粗几何对齐,无需增加额外分支或跨注意力模块即可实现高保真参考引导图像修复。

研究背景与动机

基于扩散模型的图像修复技术能够根据周围上下文生成逼真且自然的画面,但在需要忠实还原特定场景真实细节(如建筑纹理、特殊标识、特定物体外观)的任务中,传统文本引导或自监督修复模型极易产生看似合理却偏离事实的幻觉内容。为了解决这一真实性缺失问题,基于参考图像的图像修复(Reference-based Inpainting)引入了同场景或同目标的参考图像,旨在恢复缺失区域背后的客观实体而非凭空想象。

然而,如何高效且准确地利用参考图像信息面临两个紧密耦合的核心挑战:一是跨视角、尺度和光照变化下的空间对应区域定位;二是纹理细节的提取与全局语义结构的无缝融合。现有方案主要分为三类,但各有显著弊端:以 Paint-by-Example 为代表的微调方法缺乏显式空间关联,难以实现精准的内容转移;以 TransFill、CorrFill 为代表的几何对齐方法过度依赖密集特征匹配的准确性,在视角剧烈变动时容易失真;而以 BrushNet、MimicBrush、CompleteMe 为代表的架构扩展方法虽然性能优异,却额外引入了第二路 UNet 分支或复杂的跨注意力桥接模块,极大地增加了模型参数量、训练开销与部署复杂度。

针对现有方法在结构开销与利用效率之间的矛盾,本文尝试探寻一种轻量且自然的架构解法。与其引入独立的参考提取网络,不如充分挖掘预训练扩散模型自身强大的特征提取与表征兼容潜力。核心 idea:在单次前向过程中将同一个预训练扩散 UNet 复用两次(UNet-Twice),首轮提取参考图高层特征,次轮通过逐元素相加注入去噪主干;配合输入维度的图像级配对格式与区域级互补掩码构造,在粗几何对齐容忍度下实现高质量的参考引导图像合成。

方法详解

整体框架

UNet-Twice 构建于预训练 Stable Diffusion v2 修复模型之上,整体前向过程分为首通参考特征提取(UNet-I)与次通主干去噪生成(UNet-II)两个阶段。UNet-I 仅执行下采样编码路径以提取参考感知表征,UNet-II 执行完整的下采样与上采样去噪过程。由于两轮计算完全共享同一套 UNet 权重,所提取的多层级潜在特征天然处于完全对齐的表征空间与张量维度中,因此无需任何适配器或跨注意力层,仅凭逐元素相加(Element-wise Addition)即可无缝注入主干网络。在数据构造上,方法采用图像级配对保持两通输入接口的一致性,并通过特征匹配器粗对齐与互补掩码消除参考与目标上下文之间的干扰。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["目标图像 Itar 与参考图像 Iref"] --> B["阶段 1:区域级互补掩码与粗几何对齐<br/>OmniGlue 匹配 + 变换 + 互补掩码分离内外区域"]
    B --> C["阶段 2:图像级结构化输入配对<br/>组装统一 3 通道输入 I1 与 I2"]
    C --> D["阶段 3:同权双通前向与特征直连注入<br/>UNet-I 提取瓶颈特征,逐元素相加注入 UNet-II"]
    D --> E["阶段 4:多参考自注意力聚合扩展<br/>多路参考特征通道拼接并通过 Self-Attention 融合"]
    E --> F["输出:高保真修复图像"]

关键设计

1. 区域级互补掩码与粗几何对齐:消除参考与背景干扰并放宽匹配约束

为了让参考图像精准填充目标缺失区域而不破坏目标图现存的完好背景,模型在空间上构建互补的条件图。首先使用离架特征匹配器(如 OmniGlue)估计目标与参考图之间的稀疏关键点对应关系,通过单应性变换(Homography)将参考图像粗略对齐至目标坐标系得到 \(\tilde{I}_{ref}\)。随后利用训练采样掩码 \(I_{mask}\) 对目标图与对齐参考图进行空间互补调制:目标条件图保留已知背景区域,而参考条件图仅保留落入掩码空洞内部的候选内容:

\[I^{\mathrm{cond}}_{\mathrm{tar}} = I'_{tar} \odot (\mathbf{1} - I_{mask}), \quad I^{\mathrm{cond}}_{\mathrm{ref}} = \tilde{I}_{ref} \odot I_{mask}\]

其中 \(I'_{tar} = I \odot (\mathbf{1} - M)\) 为缺失区域置零的目标图像。两张条件图之和在像素空间构成粗糙但结构连续的合成先验。得益于扩散模型潜空间的纠错能力,该设计对几何对齐误差表现出极强的鲁棒性,只要语义对应内容大致落入空洞区域即可,彻底摆脱了传统几何方法对像素级精准匹配的苛刻依赖。

2. 图像级结构化输入配对:保持接口一致以避免同权网络理解歧义

标准 Stable Diffusion 修复模型接收由目标图、指示掩码与条件图像构成的拼接输入。为在不更改预训练网络接口的前提下直接复用同一权重,模型设计了结构完全对齐的输入图像组 \(\mathbf{I}_1\)\(\mathbf{I}_2\)

\[\mathbf{I}_1 = \{I_{tar}, \; \mathbf{0}, \; I_{ref}^{cond}\}, \quad \mathbf{I}_2 = \{I_{tar}, \; I_{tar}^{mask}, \; I_{tar}^{cond}\}\]

其中 \(\mathbf{I}_2\) 遵循标准修复输入,\(I_{tar}^{mask} = (\mathbf{1} - M) \odot I_{mask}\) 标识出可见背景范围。而在第一通 \(\mathbf{I}_1\) 中,原本的掩码位置被替换为全零矩阵 \(\mathbf{0}\)。这种输入设计巧妙规避了显式掩码对参考特征提取过程的误导,确保 UNet-I 在相同的张量形状下专注于提取参考图像的表征,同时使 \(\mathbf{I}_1\)\(\mathbf{I}_2\) 形成明确的图像级关联。

3. 同权双通前向与特征直连注入:无额外分支且免适配器的特征复用

以往双分支网络通常需要为参考分支单独引入数十甚至数百兆参数,并通过零卷积(Zero-Convolution)或交叉注意力层建立特征连接。UNet-Twice 让 UNet-I 与 UNet-II 完全共享权重,UNet-I 执行前向提取参考特征,并在下采样层对应位置将特征直接逐元素相加(Element-wise Addition)注入到 UNet-II。在训练过程中,UNet-I 的下采样分支与中间瓶颈层完全冻结,仅微调 UNet-II 的上采样重构层。实验表明,仅在最深处的中间瓶颈块(Mid-Block,\(n=1\))进行单一特征注入效果最佳,此时注入的是高层语义表征,既为生成提供了宏观外观指导,又避免了浅层低级特征强行覆盖导致画面出现生硬的“贴图”瑕疵。

4. 多参考自注意力聚合扩展:保持主干接口不变的多图信息融合

当存在 \(k\) 张参考图像时,框架无需重构骨干网络,而是将 UNet-I 重复运行 \(k\) 次以获取 \(k\) 组参考特征。为了使注入 UNet-II 的特征维度与单参考模式完全一致,模型将首张参考图设为主参考,并将所有 \(k\) 组潜特征在通道维度上进行拼接,输入轻量级自注意力(Self-Attention)模块进行交互与信息汇聚。注意力运算完成后,仅截取更新后的主参考特征图注入 UNet-II。该机制以最小的计算扩展代价,使多张辅助参考图像的互补视角和细节有效增强主特征表达。

损失函数 / 训练策略

UNet-Twice 继承标准隐扩散模型(LDM)的去噪评分匹配训练目标。输入图像经预训练且冻结的 VAE 编码器映射至潜在空间,在加噪步数 \(t\) 下对目标潜变量添加高斯噪声得到 \(z_t\)。模型优化仅作用于 UNet-II 的上采样参数,损失函数采用均方误差:

\[\mathcal{L} = \mathbb{E}_{\mathcal{E}(I_{tar}), \, \epsilon \sim \mathcal{N}(0, \mathbf{I}), \, t} \left[ \| \epsilon - \epsilon_\theta(z_t, t, \mathbf{I}_1, \mathbf{I}_2) \|^2 \right]\]

在具体实现中,模型采用单场景独立优化协议(Scene-specific Protocol),基于 Adam 优化器以 \(2 \times 10^{-5}\) 的学习率迭代训练 2,000 步,批大小为 8。推理阶段采用 DDIM 采样器,采样步数设为 50。

实验关键数据

主实验

评估基于 MegaDepth 数据集构建的图像对(视场重叠率 40%–70%),测试集包含 100 张未见目标图像,每张目标图配有 4 张跨视角参考图,图像分辨率统一缩放至 512×512。评测指标包括峰值信噪比(PSNR)、结构相似性(SSIM)、感知相似度(LPIPS)、DreamSim 以及 DINO 语义相似度。

单参考图像修复对比(1-Reference):

数据集 方法 PSNR ↑ SSIM ↑ LPIPS ↓ DreamSim ↓ DINO ↑
MegaDepth BrushNet 23.3474 0.9096 0.0610 0.0290 0.9568
MegaDepth Paint-by-Example (PbE) 23.3052 0.9103 0.0608 0.0311 0.9641
MegaDepth PbE + CorrFill 23.7351 0.9112 0.0594 0.0332 0.9667
MegaDepth MimicBrush 25.7770 0.9244 0.0464 0.0162 0.9798
MegaDepth CompleteMe 24.5790 0.9199 0.0506 0.0212 0.9765
MegaDepth RealFill 25.4339 0.9205 0.0500 0.0181 0.9816
MegaDepth LeftRefill 26.0812 0.9235 0.0475 0.0167 0.9807
MegaDepth UNet-Twice (本文) 27.3856 0.9319 0.0412 0.0106 0.9874

多参考图像修复对比(Multi-Reference):

参考数 方法 PSNR ↑ SSIM ↑ LPIPS ↓ DreamSim ↓ DINO ↑
2-Refs RealFill 26.2415 0.9229 0.0461 0.0123 0.9847
2-Refs LeftRefill 26.2392 0.9272 0.0423 0.0122 0.9854
2-Refs UNet-Twice (本文) 28.4582 0.9393 0.0355 0.0077 0.9915
3-Refs RealFill 26.4494 0.9250 0.0438 0.0116 0.9874
3-Refs LeftRefill 26.3340 0.9274 0.0434 0.0142 0.9858
3-Refs UNet-Twice (本文) 28.4586 0.9390 0.0355 0.0080 0.9910
4-Refs RealFill 26.6568 0.9263 0.0430 0.0113 0.9884
4-Refs LeftRefill 26.6212 0.9315 0.0455 0.0134 0.9868
4-Refs UNet-Twice (本文) 28.5596 0.9341 0.0352 0.0081 0.9912

消融实验

注入特征层数 \(n\) 的消融(\(n\) 从瓶颈层向外围浅层扩展,取值范围 \(1 \sim 13\)):

注入特征数 \(n\) PSNR ↑ SSIM ↑ LPIPS ↓ DreamSim ↓ DINO ↑ 说明
\(n=1\) (默认) 27.3858 0.9328 0.0412 0.0106 0.9874 仅注入 Mid-Block 瓶颈特征,效果最佳
\(n=2\) 27.1613 0.9311 0.0433 0.0166 0.9809 引入更浅层特征,感知误差增大
\(n=3\) 27.1112 0.9312 0.0433 0.0139 0.9814 性能小幅下滑
\(n=4\) 26.9773 0.9295 0.0440 0.0147 0.9816 持续下滑
\(n=7\) 26.9828 0.9293 0.0443 0.0147 0.9810 中浅层特征开始主导去噪
\(n=13\) 26.9106 0.9277 0.0451 0.0141 0.9811 注入全部下采样层,出现生硬贴图伪影

几何对齐方式的消融分析:

对齐配置 PSNR ↑ SSIM ↑ LPIPS ↓ DreamSim ↓ DINO ↑ 说明
Homography-base 27.3855 0.9319 0.0412 0.0106 0.9874 默认单应性变换
Homography-simple 27.4491 0.9319 0.0410 0.0105 0.9878 省略精细两阶段匹配,依然稳健
Homography-fine 27.4300 0.9318 0.0409 0.0108 0.9872 精细两阶段对齐,提升极微弱
Affine 27.2807 0.9314 0.0410 0.0106 0.9870 仿射变换(无法建模透视投影),仅微小下降

关键发现

  • 深层瓶颈特征优于多层浅层特征注入:当 \(n=1\) 时模型性能达到峰值。引入浅层特征(如 \(n=13\))时,参考图像的低级表观信息过于强硬地介入去噪过程,在对齐存在微小偏差时极易产生突兀的“复制粘贴”痕迹;而仅注入深层瓶颈特征能够让扩散模型在高级语义指导下灵活重构内容。
  • 对几何对齐精度的依赖被极大放宽:单应性粗对齐与精细两阶段对齐的各项评测指标几乎完全一致,即便使用理论上无法拟合三维透视投影的仿射变换(Affine),PSNR 也仅从 27.38 降至 27.28。这证实了隐空间扩散先验对空间错位具有出色的自愈适应能力。
  • 多参考图像增益明显:引入 2 到 4 张参考图像后,PSNR 从单参考的 27.38 dB 提升至 28.56 dB,LPIPS 降至 0.0352。自注意力模块能有效整合来自多个视角的互补未遮挡细节。

亮点与洞察

  • 零新增分支的极简架构复用:完全避免了 ControlNet 式的旁路网络构建与复杂的跨注意力适配模块,利用扩散模型自身的同权一致性直接逐元素相加,架构极其整洁高效。
  • “互补掩码 + 潜空间自愈”的设计协同:在像素级仅做粗糙的空间分离与拼贴,把复杂的细节拼接与边缘协调交由预训练扩散模型的深层隐空间表征去融合,做到了四两拨千斤。
  • 天然优雅的多参考扩展范式:通过主参考锚定与自注意力机制聚合任意数量的参考视角,保持注入维度完全固定,便于工程化部署。

局限与展望

  • 场景特异性微调的计算延迟:遵循 RealFill 的 scene-specific 评估协议,每张测试场景需要训练 2,000 步,无法实现毫秒级的纯前向零样本修复。
  • 对极端非重叠大视角的先验依赖有限:虽然放宽了精准对齐需求,但若所有参考图像均完全缺失待补区域的关键部件(重叠率过低),模型仍会退化为常规扩散模型的合理脑补。
  • 后续优化空间:未来可探索结合轻量 LoRA 或通用预训练范式,将 UNet-Twice 推向免微调(Zero-shot)的通用图像编辑管线。

相关工作与启发

  • vs RealFill: RealFill 在单场景优化中同时对目标与参考图进行全量微调,但缺乏显式的空间配对与区域互补引导;UNet-Twice 通过同权双通与互补掩码提供了明确的空间定位先验,单参考 PSNR 超出 RealFill 近 2.0 dB。
  • vs MimicBrush / CompleteMe: MimicBrush 与 CompleteMe 采用独立的第二路 UNet 提取参考信息,参数量庞大且需要专门的注意力桥接;UNet-Twice 复用同一套权重,参数开销零增加且在多项指标上均显著占优。
  • vs LeftRefill: LeftRefill 将参考图像与目标图像在画布上左右拼接(Canvas Concatenation),受限于固定的空间相对位置先验;UNet-Twice 在隐空间重合注入并施加互补掩码,空间利用更为自然直接。

评分

  • 新颖性: ⭐⭐⭐⭐ [同权双通前向与互补掩码配合,以极简设计替代沉重的双分支网络]
  • 实验充分度: ⭐⭐⭐⭐⭐ [单参考/多参考对比完整,注入深度与几何对齐消融细致深刻]
  • 写作质量: ⭐⭐⭐⭐⭐ [动机明确,问题剖析透彻,消融与图示设计逻辑严密]
  • 价值: ⭐⭐⭐⭐⭐ [为参考引导扩散生成提供了一个优雅、免额外参数且效果卓越的强基线]