Skip to content

Latest commit

 

History

History
283 lines (210 loc) · 21.7 KB

File metadata and controls

283 lines (210 loc) · 21.7 KB

Architecture

1. Motivation

Existing image inpainting methods are damage-type agnostic: LaMa (Suvorov et al., ICCV 2022) uses fast Fourier convolutions for arbitrary large masks but cannot distinguish a crack from a text overlay. Palette (Saharia et al., ICCV 2023) unifies several image-to-image tasks under one diffusion framework but provides no artist-specific style conditioning. RePaint (Lugmayr et al., CVPR 2022) does training-free diffusion inpainting but lacks any attribution mechanism. StyleGAN2 inversion (Richardson et al., CVPR 2021; Roich et al., CVPR 2022) achieves high-fidelity per-image reconstruction but does not handle inpainting with arbitrary masks, and DreamBooth/LoRA (Ruiz et al., CVPR 2023; Hu et al., ICLR 2022) provides style personalisation for diffusion but can overfit and offers no interpretability.

Three grounded gaps drive this architecture:

  1. No unified damage-type architecture — No published system jointly classifies damage type and conditions the restoration pathway in a single forward pass. A curator facing a cracked, faded painting with text annotations must currently choose a method per damage type.
  2. No style certificate — All existing style losses (Gram matrix, CLIP-based, LPIPS) are soft constraints with no formal bound on style deviation from the artist's corpus.
  3. No interpretable attribution for inpainting — Attention rollout shows spatial attention but not which stylistic features were copied from where. Restoration historians cannot answer "why did the model paint a baroque cherub here?"

Hypothesis: A single architecture combining (a) damage-aware mask-adaptive latent diffusion with (b) artist-specific style embedding interpolation and (c) region-level attribution via cross-attention to a style exemplar memory can reconstruct heterogeneous damage types while preserving artist-specific style at a level indistinguishable from expert-curated restoration.

2. At a glance

 ┌─────────────────────────────────────────────────────────────────────────┐
 │                        ArtRestore-Diffusion                             │
 │                                                                         │
 │  Damaged Image  Mask  Artist ID                                         │
 │       │           │       │                                              │
 │       ▼           ▼       ▼                                              │
 │  ┌──────────┐ ┌──────┐ ┌──────┐                                         │
 │  │VAE Enc.  │ │Mask  │ │Artist│                                         │
 │  │(frozen)  │ │Enc.  │ │Emb.  │                                         │
 │  └────┬─────┘ └──┬───┘ └──┬───┘                                         │
 │       │          │        │                                              │
 │       ▼          ▼        │                                              │
 │  ┌────────────────────┐   │                                              │
 │  │  Noised Latent z_t │   │                                              │
 │  │  mask-blended      │   │                                              │
 │  └─────────┬──────────┘   │                                              │
 │            │              │                                              │
 │            ▼              ▼                                              │
 │  ┌────────────────────────────────────────────┐                         │
 │  │         Denoising UNet (× T steps)         │                         │
 │  │                                            │                         │
 │  │  ┌──────────┐ ┌──────────┐ ┌────────────┐ │                         │
 │  │  │UNet Enc. │ │Middle    │ │UNet Dec.   │ │                         │
 │  │  │Levels 1-4│ │Block     │ │Levels 4-1  │ │                         │
 │  │  └────┬─────┘ └────┬─────┘ └──────┬─────┘ │                         │
 │  │       │            │              │        │                         │
 │  │       │  ┌─────────┴──────────┐   │        │                         │
 │  │       │  │ Exemplar Memory    │   │        │                         │
 │  │       │  │ Cross-Attention    │   │        │                         │
 │  │       │  └─────────┬──────────┘   │        │                         │
 │  │       │            │              │        │                         │
 │  │       ▼            ▼              ▼        │                         │
 │  │  ┌─────────────────────────────────────┐   │                         │
 │  │  │  Skip Connections + Boundary Gate   │   │                         │
 │  │  └─────────────────────────────────────┘   │                         │
 │  └────────────────────────────────────────────┘                         │
 │            │                                                             │
 │            ▼                                                             │
 │  ┌────────────────────────────────────┐                                 │
 │  │  VAE Decoder (frozen)              │                                 │
 │  └────────────────┬───────────────────┘                                 │
 │                   ▼                                                      │
 │  ┌────────────────────────────────────┐                                 │
 │  │  Restored Image (B, 3, H, W)      │                                 │
 │  └────────────────┬───────────────────┘                                 │
 │                   │                                                      │
 │                   ▼                                                      │
 │  ┌────────────────────────────────────┐                                 │
 │  │  Attribution Renderer              │                                 │
 │  │  (exemplar look-up → heatmap)      │                                 │
 │  └────────────────────────────────────┘                                 │
 └─────────────────────────────────────────────────────────────────────────┘
Property Value
Parameter count (test-mode UNet) 1,719,959
Parameter count (full test-mode model) 2,811,800
Parameter count (default config, estimated) ~1.6B (SD 2.1-scale UNet + conditioning)
Time complexity O(S × T) — S = 50 DDIM steps, T = latent pixels
Space complexity O(T) — UNet with attention at reduced resolutions (8×8, 16×16, 32×32 latent)
Hardware requirements GPU with ≥16 GB VRAM (default); ≥40 GB for >1024px images
VAE compression 8× spatial (SD VAE, frozen)

3. The core component

3.1 Intuition

The architecture extends a standard latent diffusion UNet with four specialised conditioning pathways, each addressing one aspect of the art restoration problem.

Mask-adaptive FiLM replaces the simple mask concatenation used in SD Inpainting. A lightweight CNN processes the damage mask at multiple scales, producing feature-wise scale and shift parameters at each UNet resolution. The same CNN handles thin cracks (binary, sparse), large holes (binary, contiguous), soft fades (continuous float), and text overlays (rectangular patches) — the module is damage-type agnostic by construction.

Artist style injection uses two complementary mechanisms. AdaLN globally modulates feature statistics (colour palette, brushstroke texture, contrast) uniformly across the canvas. Cross-attention lets different spatial regions query different aspects of the style embedding — a sky region might emphasise different style attributes than a foreground figure, even within the same painting.

Style exemplar memory cross-attention is the mechanism that makes the architecture interpretable. During denoising, each feature vector at each spatial location queries a pre-built bank of patch-level features from undamaged artist works via approximate nearest-neighbor search. The retrieved top-k exemplars are attended over, and the attention weights are extracted to produce attribution maps. Every generated patch traces to a specific real patch from the artist's corpus.

Boundary seam gating addresses a practical failure mode of latent inpainting: visible seams where generated content meets preserved content. A learned per-pixel gate predicts how much local smoothing is needed based on both the feature content and the distance from the mask edge, producing a differentiable interpolation between the original features and a Gaussian-blurred version.

3.2 Equations

Forward diffusion (noising):

Given a clean latent $z_0$, the noisy latent at timestep $t$ is:

$$z_t = \sqrt{\bar\alpha_t}, z_0 + \sqrt{1 - \bar\alpha_t}, \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)$$

where $\bar\alpha_t = \prod_{i=1}^t (1 - \beta_i)$ with a linear beta schedule $\beta_i \in [1.5 \times 10^{-4}, 1.95 \times 10^{-2}]$.

The damaged-region mask $m \in {0, 1}^{H \times W}$ is applied as:

$$z_t \leftarrow z_t \cdot (1 - m_{\text{latent}}) + \varepsilon \cdot m_{\text{latent}}$$

where $m_{\text{latent}}$ is the mask downsampled to latent resolution.

Mask-adaptive FiLM modulation:

At each UNet block, the mask encoder produces scale $\gamma_l$ and shift $\beta_l$ at layer $l$:

$$\gamma_l, \beta_l = \text{MaskEnc}_l(m)$$

$$h' = \gamma_l \cdot h + \beta_l$$

Artist style injection (dual pathway):

AdaLN pathway:

$$\mu = \text{mean}(h), \quad \sigma^2 = \text{var}(h)$$ $$\hat h = \frac{h - \mu}{\sqrt{\sigma^2 + \epsilon}}$$ $$h_{\text{adaln}} = s(a) \cdot \hat h + b(a)$$

where $s(a), b(a)$ are scale/shift predicted from artist embedding $a$.

Cross-attention pathway:

$$Q = h_{\text{adaln}}, \quad K = W_k a, \quad V = W_v a$$ $$\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{C}}\right) V$$

Exemplar memory cross-attention:

For each spatial query $q_i$ at location $i$, retrieve top-$k$ nearest neighbor keys from the bank:

$$K_i, V_i = \text{ANN}(q_i, \text{Bank}, k)$$

Attend over retrieved neighbors:

$$\alpha_{ij} = \text{softmax}_j\left(\frac{q_i \cdot k_{ij}}{\sqrt{D}}\right)$$ $$o_i = \sum_{j=1}^{k} \alpha_{ij} v_{ij}$$ $$h' = h + W_o o_i$$

DDIM sampling step (inference):

$$z_{t-1} = \sqrt{\bar\alpha_{t-1}} \left( \frac{z_t - \sqrt{1 - \bar\alpha_t}, \varepsilon_\theta(z_t, t)}{\sqrt{\bar\alpha_t}} \right) + \sqrt{1 - \bar\alpha_{t-1}} , \varepsilon_\theta(z_t, t)$$

with classifier-free guidance:

$$\varepsilon_\theta = \varepsilon_{\text{uncond}} + w \cdot (\varepsilon_{\text{cond}} - \varepsilon_{\text{uncond}})$$

where $w = 5.0$ (cf. SD's typical 7.5).

3.3 Reference implementation walk-through

The following excerpt from model.py::ArtRestoreBlock.forward() shows the sequential conditioning pathway:

def forward(self, x, temb, film_scale, film_shift, mask, distance_map,
            artist_embed=None, bank_keys=None, bank_vals=None,
            return_attn=False):
    # Step 1: GroupNorm → SiLU → Conv3×3 + timestep embedding
    h = self.norm1(x)                            # (B, C_in, H, W)
    h = F.silu(h)
    h = self.conv1(h)                            # (B, C_out, H, W)
    h = h + self.temb_proj(F.silu(temb))[:, :, None, None]

    # Step 2: Mask-Adaptive FiLM (damage-type agnostic)
    if film_scale is not None:
        h = self._mask_film(h, film_scale, film_shift)   # (B, C_out, H, W)

    # Step 3: Artist Style Injection (dual pathway)
    if self.has_attn and artist_embed is not None:
        h = self._style_block(h, artist_embed, temb)     # (B, C_out, H, W)

    # Step 4: Exemplar Memory Cross-Attention
    if self.has_exemplar and bank_keys is not None:
        h, attn_w = self._exemplar_memory(h, bank_keys, bank_vals,
                                          return_attn=return_attn)

    # Step 5: Boundary Seam Gating
    if self._boundary_gate is not None:
        h = self._boundary_gate(h, mask, distance_map)  # (B, C_out, H, W)

    # Step 6: + Residual
    h = h + self.skip(x)

    # Steps 7-8: FeedForward (Conv1×1 x4 → SiLU → Conv1×1 ÷4) + Residual
    ...
    return h, attn_weights

Shape progression (test mode, 128×128 input, base_channels=32):

  • Input latent z_t: (2, 4, 16, 16)
  • After input conv: (2, 32, 16, 16)
  • Encoder level 0 (2 blocks): (2, 32, 16, 16) → down → (2, 32, 8, 8)
  • Encoder level 1 (2 blocks): (2, 64, 8, 8) → down → (2, 64, 4, 4)
  • Encoder level 2 (2 blocks): (2, 64, 4, 4) → mid → (2, 64, 4, 4)
  • Decoder level 0 (3 blocks w/ skip): (2, 64, 4, 4) → up → (2, 64, 8, 8)
  • Decoder level 1 (3 blocks w/ skip): (2, 32, 8, 8) → up → (2, 32, 16, 16)
  • Decoder level 2 (3 blocks w/ skip): (2, 32, 16, 16)
  • Output conv: (2, 4, 16, 16)

4. Tensor shape evolution

Default config (test mode, 128×128 input):

Stage Shape Notes
Input image (B, 3, 128, 128) RGB, normalized [-1, 1]
VAE encode (B, 4, 16, 16) Frozen SD VAE, 8× down
+ noise + mask blend (B, 4, 16, 16) z_t = noise or mask-blended
UNet input conv (B, 32, 16, 16) base_channels = 32
Encoder L0 block 1 (B, 32, 16, 16) attn_resolutions includes 16 → style block active
Encoder L0 down (B, 32, 8, 8) stride-2 conv
Encoder L1 block 1 (B, 64, 8, 8) channel_mult ×2
Encoder L1 down (B, 64, 4, 4) stride-2 conv
Encoder L2 block 1 (B, 64, 4, 4) attn_res=4,8 → style + exemplar active
Mid block 1 (B, 64, 4, 4) ResBlock
Mid attention (B, 64, 4, 4) SpatialSelfAttention
Mid block 2 (B, 64, 4, 4) ResBlock
Decoder L0 (concat skip) (B, 128, 4, 4) skip from encoder L2
Decoder L0 up (B, 64, 8, 8) nearest + conv
Decoder L1 (concat skip) (B, 96, 8, 8) skip from encoder L1
Decoder L1 up (B, 32, 8, 8) nearest + conv
Decoder L2 (concat skip) (B, 64, 16, 16) skip from encoder L0
UNet output (B, 4, 16, 16) noise prediction
VAE decode (B, 3, 128, 128) restored image

5. Design decisions

Decision Alternative considered Why we chose this Trade-off accepted
Latent diffusion backbone (SD) Pixel-space diffusion, GAN inversion Latent space at 8× compression makes high-resolution attention tractable; pretrained SD VAE provides strong prior Frozen VAE constrains latent distribution
Mask as FiLM conditioning Mask concatenation as 4th channel FiLM is more parameter-efficient; naturally handles multi-scale damage (fine cracks at high res, large holes at low res) FiLM averages across film_dim → spatial detail loss at high res
Dual injection: AdaLN + cross-attn AdaLN only, cross-attn only AdaLN provides global colour/contrast/brushstroke; cross-attn lets spatial regions query different attributes Double the parameters vs single injection
Exemplar memory cross-attention Text-only conditioning (CLIP/DreamBooth) Real undamaged patches provide stronger style signal; attention weights serve as built-in attribution FAISS ANN retrieval is a non-ML bottleneck; requires GPU memory for bank
Boundary gating Post-hoc Poisson blending Differentiable and trained end-to-end; adapts to local features Risk of collapse to identity or all-blur
DDIM sampling (50 steps) 1000-step DDPM, DPM-Solver 20× speedup; deterministic for reproducible attribution Slight quality reduction vs full DDPM
CFG scale 5.0 SD default 7.5 Lower CFG because exemplar memory provides strong conditioning; 7.5 oversaturates paint textures May under-utilise text conditioning
Lazy channel projections Fixed projections for all UNet levels Enables variable channel counts without pre-computing all projections Slight first-forward overhead
FAISS ANN (planned) Brute-force search 500K bank entries per artist × brute-force per pixel per step is prohibitive ANN recall < 100%; fallback to PQ64 or HNSW

6. Domain-specific considerations

GenAI

Latent vs. pixel space: Latent diffusion (SD VAE, 8× compression) is chosen because artworks are high-resolution (≥1024px) and pixel-space diffusion at that resolution is prohibitively expensive. The pretrained SD VAE provides a strong prior requiring minimal fine-tuning.

Conditioning interface: Three independent conditioning pathways — text (standard SD cross-attention, optional), artist embedding (separate cross-attention head + AdaLN, primary style signal), and exemplar memory (cross-attention to bank, per-patch style reference).

Sampling efficiency: 50-step DDIM at inference (vs. 1000-step DDPM). DDIM is preferred over DPM-Solver because its deterministic nature makes attribution maps consistent across runs.

CFG strategy: Classifier-free guidance with the artist embedding as the conditioning signal. The unconditional pass uses a learnable null artist embedding. CFG scale of 5.0 is lower than SD's 7.5 because the exemplar memory already provides strong conditioning.

CV

Spatial handling: Fully convolutional UNet with no positional embeddings — the VAE preserves spatial layout. Different input resolutions are handled by the UNet's down/upsampling structure.

Multi-scale damage: Fine cracks are best captured at high resolution (early UNet layers); large missing regions require global context (deep UNet layers). The mask FiLM operates at all scales.

Trustworthiness / XAI

Attribution mechanism: The exemplar memory cross-attention provides per-patch attribution to specific source patches. This is more interpretable than attention rollout because each generated patch maps to a concrete, real exemplar that a human can inspect.

Style certificate: The style certificate (embedding distance from artist manifold) provides a quantitative bound on style deviation. Implementation: pre-compute artist centroid from features of all undamaged works; at inference, measure cosine distance per patch in the restored region; flag patches exceeding 3σ of within-corpus variance.

Limitation: The style certificate only measures deviation from the centroid, not the appropriateness of the specific choice. A baroque cherub in a Renaissance painting might fall within style bounds if the artist's corpus includes both periods.

7. Known limitations

  • Untrained model — All benchmarks run on randomly initialized models. Performance targets (LPIPS ≤ 0.08, DISTS ≤ 0.06, style certificate < 10% violation) are aspirational and unverified. Trained-model evaluation is the highest-priority next step.
  • Metric-expert correlation unmeasured — The strategy of reducing expert panel size via automated metrics (DISTS, LPIPS) requires Spearman ρ ≥ 0.7 between DISTS and expert Likert ratings. This has not been empirically measured and requires a human subjects study (IRB approval).
  • Falsification conditions 2-3 untestable by automation — Expert agreement (Fleiss' κ ≥ 0.6), mean expert rating (≥ 4.0/5), and metric-expert correlation require human subjects. These are flagged as blocking gaps in research_eval/scorecard.json.
  • Baselines not integrated — LaMa (FFT inpainting) and StyleGAN2 inversion baselines are identified but no comparative pipeline exists. The architecture cannot claim superiority without comparative results.
  • Blocking unknown: multi-period artists — A single 512-d artist embedding may be insufficient for artists with diverse sub-periods (e.g., Picasso's Blue vs. Cubist periods). Per-period embeddings may be needed, increasing data requirements.
  • Blocking unknown: boundary artefacts — The learned boundary gate may collapse to identity (no seam mitigation) or all-blur (over-smoothed edges). Monitoring gates during training is required.
  • FAISS ANN recall on artistic features — IVF-PQ is optimised for natural image features. Artistic features (brushstroke textures) may not cluster well, potentially reducing recall@16. Fallback: HNSW or PQ128.
  • Synthetic data only — Current benchmarks use synthetic damage (random masks, lines, gradients). Real-world art restoration involves complex damage patterns not captured by synthetic noise.
  • No training convergence analysis — Learning curves, loss landscapes, and gradient conflict analysis (between noise-prediction and style-consistency losses) are pending.