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

Conference: ECCV 2026
Paper: ECCV Official
Area: Image Generation
Keywords: reference-based inpainting / diffusion models / weight sharing / complementary masking / image restoration

TL;DR

UNet-Twice proposes a simple yet effective framework that reuses the exact same diffusion UNet in two passes per forward step alongside complementary masking and coarse geometric alignment, achieving state-of-the-art reference-based inpainting without any extra branches or cross-attention bridges.

Background & Motivation

Diffusion models have achieved exceptional generative realism in image inpainting, synthesizing visually plausible structures to complete missing regions based on ambient context. However, in applications that demand faithful recovery of specific scenes or objects—such as restoring architectural structures, proprietary textures, or distinct landmarks—standard text-guided or self-supervised inpainting networks frequently hallucinate plausible yet factually incorrect content. Reference-based inpainting addresses this critical limitation by supplying additional reference views of the same scene, aiming to recover authentic visual content rather than imaginative substitutes.

Effectively harvesting reference guidance entails two tightly coupled challenges: accurately identifying spatial correspondences across viewpoints, scales, and illumination conditions, and seamlessly transferring fine appearance details without disrupting the unmasked target context. Prior research typically falls into three paradigms, all bearing notable drawbacks. Straightforward fine-tuning methods (such as Paint-by-Example) lack explicit spatial associations and thus struggle with accurate local feature transfer. Geometry-alignment approaches (such as TransFill and CorrFill) depend heavily on precise pixel-level dense correspondences, which degrades severely under significant viewpoint shifts. Meanwhile, architecture-augmented frameworks (such as BrushNet, MimicBrush, and CompleteMe) introduce an entire second UNet branch and complex cross-attention bridges, incurring substantial parameter overhead, training burden, and deployment friction.

This paper tackles the trade-off between architectural complexity and reference utilization efficiency by unlocking the latent representational capacity already present in pretrained diffusion models. Rather than adding auxiliary networks, the model capitalizes on the inherent compatibility of identical feature representations. Core idea: reuse the exact same pretrained diffusion UNet twice within a single forward pass (UNet-Twice) to extract reference latent features and directly inject them into the denoising backbone via element-wise addition, combined with an image-level paired interface and region-level complementary masking that robustly tolerates coarse geometric misalignment.

Method

Overall Architecture

UNet-Twice builds upon a pretrained Stable Diffusion v2 inpainting model. The overall forward execution is organized into two consecutive passes: the reference extraction pass (UNet-I) and the target denoising pass (UNet-II). UNet-I runs exclusively through the down-sampling pathway to capture multi-scale reference-aware latent features, while UNet-II executes the full down-sampling, bottleneck, and up-sampling denoising trajectory. Because both passes are parameterized by identical network weights, their representations naturally align in channel dimensions and feature space, enabling direct injection into UNet-II via element-wise addition without adapters or cross-attention layers. At the data interface, structured input pairing preserves format consistency across passes, while keypoint-based coarse warping and complementary masking eliminate interference between reference cues and visible target context.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Target Image Itar and Reference Image Iref"] --> B["Stage 1: Region-level Complementary Masking and Alignment<br/>OmniGlue matching + homography warp + complementary split"]
    B --> C["Stage 2: Image-level Paired Input Construction<br/>Construct unified 3-tensor inputs I1 and I2"]
    C --> D["Stage 3: Shared-Weight Dual-Pass Forward and Direct Injection<br/>UNet-I extracts bottleneck features, injected into UNet-II via addition"]
    D --> E["Stage 4: Multi-Reference Self-Attention Aggregation<br/>Concatenate k reference features and fuse via self-attention"]
    E --> F["Output: High-fidelity Inpainted Image"]

Key Designs

1. Region-level Complementary Masking and Coarse Alignment: Eliminating Spatial Interference Under Relaxed Matching Constraints

To ensure the reference image supplies candidate details strictly for the masked target void without contaminating valid target backgrounds, the framework enforces spatial complementarity between both inputs. An off-the-shelf feature matcher (OmniGlue) first detects sparse correspondences to estimate a coarse homography warp, producing an aligned reference image \(\tilde{I}_{ref}\). Next, a random binary mask \(I_{mask}\) spatially modulates both the target and aligned reference frames such that the target provides valid context outside the hole while the reference provides candidate content strictly inside:

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

where \(I'_{tar} = I \odot (\mathbf{1} - M)\) denotes the target image with masked regions zeroed out. Summing both conditions forms a composite prior in image space. Because latent diffusion representations possess strong error-correction capabilities, this design eliminates the need for sub-pixel registration; as long as the warped reference content roughly falls within the masked target area, the network reliably synthesizes seamless boundaries.

2. Image-level Paired Input Construction: Preserving Interface Consistency to Prevent Weight-Sharing Ambiguity

Standard Stable Diffusion inpainting requires a composite input consisting of the noisy target image, an indicator mask, and a conditioning background image. To reuse identical weights across both passes without confusing the network, the framework defines structurally symmetric input sets \(\mathbf{I}_1\) and \(\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}\}\]

where \(\mathbf{I}_2\) adheres to the standard inpainting formulation with \(I_{tar}^{mask} = (\mathbf{1} - M) \odot I_{mask}\) indicating preserved context. In the auxiliary reference pass \(\mathbf{I}_1\), the indicator mask is replaced by an all-zero tensor \(\mathbf{0}\). This suppresses false explicit mask priors during the first pass while preserving exact tensor shapes, ensuring that UNet-I cleanly extracts appearance features driven by \(I_{ref}^{cond}\) and establishing a symmetric image-level pairing across the two passes.

3. Shared-Weight Dual-Pass Forward and Direct Injection: Zero-Parameter Feature Reuse Without Auxiliary Adapters

Prior dual-branch systems instantiate dedicated reference encoders containing dozens or hundreds of megabytes of extra parameters, requiring zero-convolutions or cross-attention modules for feature injection. UNet-Twice reuses the exact same backbone weights for both passes. During training, the down-sampling pathway and mid-block of UNet-I and UNet-II are completely frozen, and only the up-sampling layers of UNet-II are optimized. The reference features from UNet-I are directly added element-wise into the corresponding layers of UNet-II. Systematic ablations reveal that injecting exclusively at the deepest mid-block bottleneck (\(n=1\)) delivers optimal performance; high-level semantic features guide plausible contextual reconstruction without letting raw low-level textures induce rigid, artificial paste-in artifacts.

4. Multi-Reference Self-Attention Aggregation: Seamless Multi-View Fusion Under an Unchanged Injection Interface

When \(k\) reference images are available, the framework executes UNet-I \(k\) times to produce \(k\) reference latent feature tensors without altering the primary denoising backbone. To keep the injected tensor dimensions identical to the single-reference setting, the first reference image is designated as the primary anchor. All \(k\) feature maps are concatenated along the channel dimension and processed by a lightweight self-attention layer to facilitate cross-view information exchange. Only the updated primary reference feature is subsequently retained and injected into UNet-II. This allows complementary viewpoints to enrich the guidance signal while preserving the original element-wise addition interface.

Loss & Training

UNet-Twice adheres to the standard latent diffusion denoising score-matching objective. Input images are mapped into the latent domain using the frozen pretrained VAE encoder \(\mathcal{E}(\cdot)\). At diffusion step \(t\), Gaussian noise \(\epsilon \sim \mathcal{N}(0, \mathbf{I})\) is added to the target latent to yield \(z_t\). Optimization is restricted entirely to the up-sampling layers of UNet-II using the mean squared error loss:

\[\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]\]

Following the scene-specific training protocol established by RealFill, the model is fine-tuned for 2,000 iterations per scene with a batch size of 8 using the Adam optimizer at a learning rate of \(2 \times 10^{-5}\). Inference uses the DDIM sampler across 50 denoising steps.

Key Experimental Results

Main Results

Experiments are conducted on MegaDepth image pairs exhibiting 40%–70% field-of-view overlap. The evaluation set consists of 100 disjoint target images, each paired with four multi-view reference images, resized to 512×512 resolution. Evaluations report Peak Signal-to-Noise Ratio (PSNR), Structural Similarity (SSIM), Learned Perceptual Image Patch Similarity (LPIPS), DreamSim, and DINO visual similarity.

Single-Reference Inpainting Benchmark (1-Reference):

Dataset	Method	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 (Ours)	27.3856	0.9319	0.0412	0.0106	0.9874

Multi-Reference Inpainting Benchmark (Multi-Reference):

Setting	Method	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 (Ours)	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 (Ours)	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 (Ours)	28.5596	0.9341	0.0352	0.0081	0.9912

Ablation Study

Impact of the Number of Injected Features \(n\) (expanding outward from the mid-block bottleneck, \(n \in \{1, \dots, 13\}\)):

Injected Features \(n\)	PSNR ↑	SSIM ↑	LPIPS ↓	DreamSim ↓	DINO ↑	Note
\(n=1\) (Default)	27.3858	0.9328	0.0412	0.0106	0.9874	Deepest mid-block feature only; optimal balance
\(n=2\)	27.1613	0.9311	0.0433	0.0166	0.9809	Adding shallower features raises perceptual error
\(n=3\)	27.1112	0.9312	0.0433	0.0139	0.9814	Slight degradation
\(n=4\)	26.9773	0.9295	0.0440	0.0147	0.9816	Consistent drop
\(n=7\)	26.9828	0.9293	0.0443	0.0147	0.9810	Low-level textures begin dominating
\(n=13\)	26.9106	0.9277	0.0451	0.0141	0.9811	Full down-sampling injection; rigid copy-paste artifacts

Ablation on Geometric Alignment Strategies:

Alignment Strategy	PSNR ↑	SSIM ↑	LPIPS ↓	DreamSim ↓	DINO ↑	Note
Homography-base	27.3855	0.9319	0.0412	0.0106	0.9874	Standard coarse homography
Homography-simple	27.4491	0.9319	0.0410	0.0105	0.9878	Omits secondary fine matching; remains strong
Homography-fine	27.4300	0.9318	0.0409	0.0108	0.9872	Fine two-stage alignment; negligible gain
Affine	27.2807	0.9314	0.0410	0.0106	0.9870	Affine warp (cannot model perspective); marginal drop

Key Findings

• Deep Bottleneck Injection Outperforms Shallow Multiscale Injection: Peak fidelity is achieved at \(n=1\). Progressively injecting shallower features (\(n=13\)) causes low-level appearance patterns to dominate the denoising trajectory, producing jarring paste-in visual artifacts whenever geometric alignment is slightly off. Deep bottleneck features provide semantic appearance guidance while letting the generative backbone synthesize harmonized details.
• Robustness to Imperfect Geometric Alignment: Differences across homography refinement variants are virtually negligible. Even under affine transformations that fundamentally fail to model perspective projections, PSNR drops by only 0.1 dB (27.28 vs. 27.38). This confirms that diffusion latent representations provide exceptional error-correction tolerance, lifting the requirement for precise pixel-level registration.
• Consistent Multi-Reference Scaling: Scaling from 1 to 4 reference views elevates PSNR from 27.38 dB to 28.56 dB and lowers LPIPS to 0.0352, demonstrating that the self-attention aggregation module successfully extracts and merges complementary non-occluded cues across perspectives.

Highlights & Insights

• Zero-Parameter Architecture Reuse: Completely bypasses costly auxiliary reference branches and dedicated attention bridging layers, using the identical UNet weights via element-wise addition for maximum simplicity and efficiency.
• Synergy of Complementary Masking and Latent Self-Correction: Implements simple coarse masking in image space while delegating complex edge blending and structural harmonization to the generative latent space of the diffusion backbone.
• Scalable Multi-Reference Formulation: Anchoring on a primary reference and fusing extra views via self-attention maintains an invariant injection interface, rendering multi-image extensions remarkably clean.

Limitations & Future Work

• Scene-Specific Fine-Tuning Overhead: Adhering to the RealFill evaluation protocol necessitates 2,000 optimization steps per test scene, which is computationally expensive and unsuitable for real-time inference.
• Degradation Under Zero-Overlap Viewpoints: While forgiving of spatial misalignment, the framework still requires the reference image to contain the missing object; if content is entirely absent across all views, generation reverts to unguided context hallucination.
• Future Directions: Exploring low-rank adaptation (LoRA) or large-scale pretraining to adapt UNet-Twice into a feed-forward, zero-shot reference inpainting framework.

Related Work & Insights

• vs RealFill: RealFill fine-tunes all parameters on unmasked target and reference pairs without explicit spatial structure; UNet-Twice incorporates paired inputs and complementary masking to deliver structured spatial priors, outperforming RealFill by nearly 2.0 dB in PSNR under single-reference settings.
• vs MimicBrush / CompleteMe: MimicBrush and CompleteMe introduce dedicated secondary UNets with heavy parameter footprints and complex attention modules; UNet-Twice reuses the same backbone with zero added branch parameters and superior quantitative scores.
• vs LeftRefill: LeftRefill concatenates reference and target images horizontally along the canvas, imposing rigid spatial bias; UNet-Twice overlays features in the latent domain with complementary masks, enabling more natural spatial integration.

Rating

• Novelty: ⭐⭐⭐⭐ [Reusing the identical UNet twice with complementary masking elegantly eliminates auxiliary branches]
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ [Extensive single/multi-reference benchmarks and thorough injection depth/alignment ablations]
• Writing Quality: ⭐⭐⭐⭐⭐ [Clear motivation, structured exposition, and well-designed experimental figures]
• Value: ⭐⭐⭐⭐⭐ [Provides an elegant, zero-added-parameter reference-guided inpainting baseline for diffusion systems]