Leveraging the Powerful Attention of a Pre-trained Diffusion Model for Exemplar-based Image Colorization

TL;DR

This work tackles exemplar-based image colorization by leveraging the self-attention of a pre-trained diffusion model to align grayscale inputs with a reference color image. It introduces a fine-tuning-free framework with dual attention-guided color transfer and classifier-free colorization guidance, enabling semantically aware color transfer without retraining. The method achieves state-of-the-art performance on 335 input–reference pairs (FID $=$ $95.27$, SI-FID $=$ $5.51$) and demonstrates robustness in cross-domain and exemplar-free comparisons. Overall, it showcases the practical potential of reusing diffusion-model attention for semantically guided colorization, with broad applicability in restoration, archiving, and creative media.

Abstract

Exemplar-based image colorization aims to colorize a grayscale image using a reference color image, ensuring that reference colors are applied to corresponding input regions based on their semantic similarity. To achieve accurate semantic matching between regions, we leverage the self-attention module of a pre-trained diffusion model, which is trained on a large dataset and exhibits powerful attention capabilities. To harness this power, we propose a novel, fine-tuning-free approach based on a pre-trained diffusion model, making two key contributions. First, we introduce dual attention-guided color transfer. We utilize the self-attention module to compute an attention map between the input and reference images, effectively capturing semantic correspondences. The color features from the reference image is then transferred to the semantically matching regions of the input image, guided by this attention map, and finally, the grayscale features are replaced with the corresponding color features. Notably, we utilize dual attention to calculate attention maps separately for the grayscale and color images, achieving more precise semantic alignment. Second, we propose classifier-free colorization guidance, which enhances the transferred colors by combining color-transferred and non-color-transferred outputs. This process improves the quality of colorization. Our experimental results demonstrate that our method outperforms existing techniques in terms of image quality and fidelity to the reference. Specifically, we use 335 input-reference pairs from previous research, achieving an FID of 95.27 (image quality) and an SI-FID of 5.51 (fidelity to the reference). Our source code is available at https://github.com/satoshi-kosugi/powerful-attention.

Figures (13)

• Figure 1: Overview of our method. First, $z_0^{\rm in}$, $z_0^{\rm ref}$, and $z_0^{\rm refL}$ are given; these are the latent variables corresponding to the input, reference, and grayscale reference, respectively. We apply DDIM inversion to these latent variables to obtain $z_T^{\rm in}$, $z_T^{\rm ref}$, and $z_T^{\rm refL}$. Next, we set $z_T^{\rm in}$ as the initial value for $z_T^{\rm out}$ and apply the colorization process to $z_T^{\rm out}$. This process consists of two key components: 1) dual attention-guided color transfer and 2) classifier-free colorization guidance. Finally, we obtain the colorized output $z_0^{\rm out}$.
• Figure 2: Illustration of (a) general usage of self-attention rombach2022high and (b) dual attention-guided color transfer.
• Figure 3: Illustration of classifier-free colorization guidance.
• Figure 4: Qualitative comparison.
• Figure 5: Qualitative comparison in ablation studies.