SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peiqing Yang $^{1}$ $^{1}$ S-Lab, Nanyang Technological University  
 $^{2}$ SenseTime Research, Singapore  
https://github.com/pq-yang/PGDiff
Exploiting pre-trained diffusion models for restoration has recently become a favored alternative to the traditional task-specific training approach. Previous works have achieved noteworthy success by limiting the solution space using explicit degradation models. However, these methods often fall short when faced with complex degradations as they generally cannot be precisely modeled. In this paper, we propose PGDiff by introducing partial guidance, a fresh perspective that is more adaptable to real-world degradations compared to existing works. Rather than specifically defining the degradation process, our approach models the desired properties, such as image structure and color statistics of high-quality images, and applies this guidance during the reverse diffusion process. These properties are readily available and make no assumptions about the degradation process. When combined with a diffusion prior, this partial guidance can deliver appealing results across a range of restoration tasks. Additionally, PGDiff can be extended to handle composite tasks by consolidating multiple high-quality image properties, achieved by integrating the guidance from respective tasks. Experimental results demonstrate that our method not only outperforms existing diffusion-prior-based approaches but also competes favorably with task-specific models.
Recent years have seen diffusion models achieve outstanding results in synthesizing realistic details across various content [33, 29, 7, 15, 34]. The rich generative prior inherent in these models opens up a vast array of possibilities for tasks like super-resolution, inpainting, and colorization. Consequently, there has been a growing interest in formulating efficient guidance strategies for pre-trained diffusion models, enabling their successful adaptation to various restoration tasks [9, 42, 20, 37].
A common approach [9, 42, 20] is to constrain the solution space of intermediate outputs during the denoising process<sup>1</sup>. At each iteration, the intermediate output is modified such that its degraded counterpart is guided towards the input low-quality (LQ) image. Existing works achieve this goal either by using a closed-form solution [42, 20] or back-propagating simple losses [9]. These methods are versatile in the sense that the pre-trained diffusion model can be adapted to various tasks without fine-tuning, as long as the degradation process is known in advance.
While possessing great versatility, the aforementioned methods are inevitably limited in generalizability due to the need for prior knowledge of the degradation process. In particular, a closed-form solution generally does not exist except for special cases such as linear operators. In addition, back-propagating losses demand differentiability of the degradation process, which is violated for many degradations such as JPEG compression. Importantly, degradations in the wild often consist
It refers to the reverse diffusion process, not the image denoising task.
Figure 1: Overview of Our PGDiff Framework for Versatile Face Restoration. Here, we take the colorization task as an example to illustrate our inference pipeline. One may refer to Table 1 for the corresponding details (e.g., property, classifier, and target) of other tasks. We show that our method can handle a wide range of tasks, including (a) blind face restoration, (b) face colorization, (c) face inpainting, and also composite tasks such as (d) old photo restoration.
of a mixture of degradations [41], and hence, it is difficult, if not impossible, to model them accurately. As a result, existing works generally limit the scope to simplified cases, such as fixed-kernel downsampling. The generalization to real-world degradations remains a formidable challenge.
Motivated by the above, instead of modeling the degradation process, we propose to model the desired properties of high-quality (HQ) images. The merit of such guidance is the agnosticity to the degradation process. However, it remains unclear what properties are desired and how appropriate guidance can be constructed. Through our extensive experiments, we find that with diffusion prior acting as a natural image regularization, one could simply guide the denoising process with easily accessible properties, such as image structure and color statistics. For example, as shown in Fig. 1, one could generate plausible outputs simply by providing guidance on the lightness and the statistics (i.e., mean and variance) of each color channel, without knowing the exact decolorization process. By constraining the HQ image space, our idea bypasses the difficulty of knowing the prior relation between LQ and HQ images, thus improving generalizability.
In this work, we devise a simple yet effective instantiation named PGDiff by introducing partial guidance. PGDiff adopts classifier guidance [7] to constrain the denoising process. Each image property corresponds to a classifier, and the intermediate outputs are updated by back-propagating the gradient computed on the loss between the classifier output and the target property. Since our partial guidance is agnostic to the degradation process, it can be easily extended to complex tasks by compositing multiple properties. For instance, the task of old photo restoration can be regarded as a combination of restoration, inpainting, and colorization, and the resultant guidance is represented as a weighted sum of the guidance in the respective task. We also demonstrate that common losses such as perceptual loss [2, 16] and adversarial loss [22] can be incorporated for further performance gain.
Contributions. Our main contributions include i) a new concept of adapting diffusion models to restoration without presumptions of the degradation process. We show that it suffices to guide the denoising process with easily accessible properties in the HQ image space, with diffusion prior acting as regularization, and ii) partial guidance, a versatile approach that is applicable to a broad range of image restoration and enhancement tasks. Furthermore, it allows flexible combinations of guidance for intricate tasks. We conduct extensive experiments to demonstrate the effectiveness of PGDiff on a variety of challenging tasks including blind face restoration and old photo restoration. We also demonstrate interesting applications, such as reference-based restoration. The results confirm the superiority of PGDiff over previous state-of-the-art methods.
Generative Prior for Restoration. Generative prior has been widely adopted for a range of image restoration tasks, including super-resolution, inpainting, and colorization. One prominent approach in
this field is the use of pre-trained generative adversarial networks (GANs) [10, 18, 1]. For instance, GAN-inversion [26, 11, 28] inverts a corrupted image to a latent code, which is then used for generating a clean image. Another direction is to incorporate the prior into an encoder-decoder architecture [3, 4, 40, 44], bypassing the lengthy optimization during inference. VQVAE [35] is also commonly used as generative prior. Existing works [48, 13, 43, 47] generally first train a VQVAE with a reconstruction objective, followed by a fine-tuning stage to adapt to the subsequent restoration task. Recently, diffusion models have gained increasing attention due to their unprecedented performance in various generation tasks [33, 29, 7, 15, 34], and such attention has led to interest in leveraging them as a prior for restoration.
Diffusion Prior. There has been a growing interest in formulating efficient guidance strategies for pre-trained diffusion models, enabling their successful adaptation to various restoration tasks [9, 42, 20, 37, 39]. Among them, DDRM [20], DDNM [42], and GDP [9] adopt a zero-shot approach to adapt a pre-trained diffusion model for restoration without the need of task-specific training. At each iteration, the intermediate output is modified such that its degraded counterpart is guided towards the input low-quality image. This is achieved under an assumed degradation process, either in the form of a fixed linear matrix [42, 20] or a parameterized degradation model [9], with learnable parameters representing degradation extents. In this work, we also exploit the generative prior of a pre-trained diffusion model by formulating efficient guidance for it, but unlike existing works that limit the solution space using explicit degradations [9, 42, 20], we propose to model the desired properties of high-quality images. Such design is agnostic to the degradation process, circumventing the difficulty of modeling the degradation process.
PGDiff is based on diffusion models. In this section, we first introduce the background related to our method in Sec. 3.1, and the details of our method are presented in Sec. 3.2.
Diffusion Models. The diffusion model [33] is a class of generative models that learn to model a data distribution  $p(x)$ . In particular, the forward process is a process that iteratively adds Gaussian noise to an input  $x_0 \sim p(x)$ , and the reverse process progressively converts the data from the noise distribution back to the data distribution, often known as the denoising process.
For an unconditional diffusion model with  $T$  discrete steps, at each step  $t$ , there exists a transition distribution  $q(x_{t}|x_{t - 1})$  with variance schedule  $\beta_{t}$  [15]:
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(x _ {t}; \sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} \mathbf {I}\right). \tag {1}
$$
Under the reparameterization trick,  $x_{t}$  can be written as:
$$
x _ {t} = \sqrt {\alpha_ {t}} x _ {t - 1} + \sqrt {1 - \alpha_ {t}} \epsilon , \tag {2}
$$
where  $\alpha_{t} = 1 - \beta_{t}$  and  $\epsilon \sim \mathcal{N}(\epsilon ;\mathbf{0},\mathbf{I})$ . Recursively, let  $\bar{\alpha}_{t} = \prod_{i = 1}^{t}\alpha_{i}$ , we have
$$
x _ {t} = \sqrt {\bar {\alpha} _ {t}} x _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon . \tag {3}
$$
During sampling, the process starts with a pure Gaussian noise  $x_{T} \sim \mathcal{N}(x_{T}; \mathbf{0}, \mathbf{I})$  and iteratively performs the denoising step. In practice, the ground-truth denoising step is approximated [7] by  $p_{\theta}(x_{t-1}|x_t)$  as:
$$
p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {N} \left(\mu_ {\theta} \left(x _ {t}, t\right), \Sigma_ {\theta} \left(x _ {t}, t\right)\right), \tag {4}
$$
where  $\Sigma_{\theta}(x_t,t)$  is a constant depending on pre-defined  $\beta_{t}$ , and  $\mu_{\theta}(x_t,t)$  is generally parameterized by a network  $\epsilon_{\theta}(x_t,t)$ :
$$
\mu_ {\theta} \left(x _ {t}, t\right) = \frac {1}{\sqrt {\bar {\alpha} _ {t}}} \left(x _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \epsilon_ {\theta} \left(x _ {t}, t\right)\right). \tag {5}
$$
From Eq. (3), one can also directly approximate  $x_0$  from  $\epsilon_{\theta}$ :
$$
\hat {x} _ {0} = \frac {1}{\sqrt {\bar {\alpha} _ {t}}} x _ {t} - \sqrt {\frac {1 - \bar {\alpha} _ {t}}{\bar {\alpha} _ {t}}} \epsilon_ {\theta} (x _ {t}, t). \tag {6}
$$
Algorithm 1 Sampling with partial guidance, given a diffusion model  $(\mu_{\theta}(x_t,t),\Sigma_{\theta}(x_t,t))$  , classifier  $p_{\phi}(y|\hat{x}_0)$  , target  $y$  gradient scale  $s$  , range for multiple gradient steps  $S_{start}$  and  $S_{end}$  , and the number of gradient steps  $N$  . The dynamic guidance weight  $s_{norm}$  is omitted for brevity.
Input: a low-quality image  $y_{0}$ $x_{T} \gets$  sample from  $\mathcal{N}(\mathbf{0},\mathbf{I})$   
for  $t$  from  $T$  to 1 do  
 $\mu, \Sigma \gets \mu_{\theta}(x_{t},t), \Sigma_{\theta}(x_{t},t)$ $\hat{x}_{0} \gets \frac{1}{\sqrt{\bar{\alpha}_{t}}} x_{t} - \sqrt{\frac{1 - \bar{\alpha}_{t}}{\bar{\alpha}_{t}}} \epsilon_{\theta}(x_{t},t)$   
if  $S_{start} \leq t \leq S_{end}$  then  
repeat  
 $x_{t} \gets$  sample from  $\mathcal{N}(\mu - s\Sigma \nabla_{\hat{x}_{0}} \| y - p_{\phi}(y|\hat{x}_{0}) \|_{2}^{2}, \Sigma)$ $\hat{x}_{0} \gets \frac{1}{\sqrt{\bar{\alpha}_{t}}} x_{t} - \sqrt{\frac{1 - \bar{\alpha}_{t}}{\bar{\alpha}_{t}}} \epsilon_{\theta}(x_{t},t)$   
until  $N - 1$  times  
end if  
 $x_{t-1} \gets$  sample from  $\mathcal{N}(\mu - s\Sigma \nabla_{\hat{x}_{0}} \| y - p_{\phi}(y|\hat{x}_{0}) \|_{2}^{2}, \Sigma)$   
end for  
return  $x_{0}$
Table 1: Examples of Partial Guidance. Each image property corresponds to a classifier, and each task involves one or multiple properties as guidance. The target value of each property is generally obtained either from the input image  $y_0$  or the denoised intermediate output  $\hat{x}_0$ . For a composite task, we simply decompose it into multiple tasks and combine the respective guidance. Here, Clean denotes a pre-trained restorer detailed in Sec. 4.1, Identity refers to an identity mapping, and  $y_{ref}$  represents a reference image containing entity with the same identity as  $y_0$ .
Task  $\equiv$  Property  $\equiv$  Target:  $y$  Classfier:  $p_{\phi}(y|\hat{x}_0)$    
Homogeneous Inpainting Unmasked Region  $\mathbf{M}$  ask(yo) Mask   
Colorization Lightness RGB2gray(yo) rgb2gray   
Color Statistics Color Statistics AdaIN(xo）[18] Identity   
Restoration Smooth Semantics Clean(yo) Identity   
Ref-Based Restoration Smooth Semantics Clean(yo) Identity   
Identity Reference ArcFace(yref) ArcFace [6]   
Task Composition   
Composite Old Photo Restoration (w/ scratches) Restoration + Inpainting + Colorization
Classifier Guidance. Classifier guidance is used to guide an unconditional diffusion model so that conditional generation is achieved. Let  $y$  be the target and  $p_{\phi}(y|x)$  be a classifier, the conditional distribution is approximated as a Gaussian similar to the unconditional counterpart, but with the mean shifted by  $\Sigma_{\theta}(x_t,t)g$  [7]:
$$
p _ {\theta , \phi} \left(x _ {t - 1} \mid x _ {t}, y\right) \approx \mathcal {N} \left(\mu_ {\theta} \left(x _ {t}, t\right) + \Sigma_ {\theta} \left(x _ {t}, t\right) g, \Sigma_ {\theta} \left(x _ {t}, t\right)\right), \tag {7}
$$
where  $g = \nabla_{x}\log p_{\phi}(y|x)|_{x = \mu_{\theta}(x_t,t)}$ . The gradient  $g$  acts as a guidance that leads the unconditional sampling distribution towards the condition target  $y$ .
3.2 Partial Guidance
Our partial guidance does not assume any prior knowledge of the degradation process. Instead, with diffusion prior acting as a regularization, we provide guidance only on the desired properties of high-quality images. The key to PGDiff is to construct proper guidance for each task. In this section, we will discuss the overall framework, and the formulation of the guidance for each task is presented in Sec. 4. The overview is summarized in Fig. 1 and Algorithm 1.
Property and Classifier. The first step of PGDiff is to determine the desired properties which the high-quality output possesses. As summarized in Table 1, each image property corresponds to a classifier  $p_{\phi}(y|\hat{x}_0)$ , and the intermediate outputs  $x_{t}$  are updated by back-propagating the gradient computed on the loss between the classifier output and the target  $y$ .
Input
Figure 2: Comparison on Blind Face Restoration. Input faces are corrupted by real-world degradations. Our PGDiff produces high-quality faces with faithful details. (Zoom in for best view.)
PULSE
GFP-GAN
CodeFormer
GDP
DDNM
DifFace
Given a specific property (e.g., lightness), we construct the corresponding classifier (e.g., rgb2gray), and apply classifier guidance during the reverse diffusion process as shown in Fig. 1. Although our PGDiff is conceptually similar to classifier guidance, we find that the conventional guidance scheme often leads to suboptimal performance. In this work, we borrow ideas from existing works [37, 5] and adopt a dynamic guidance scheme, which introduces adjustments to the guidance weight and number of gradient steps for enhanced quality and controllability.
Dynamic Guidance Scheme. Our dynamic guidance scheme consists of two components. First, we observe that the conventional classifier guidance, which adopts a constant gradient scale  $s$ , often fails in guiding the output towards the target value. This is especially unfavourable in tasks where high similarity to the target is desired, such as inpainting and colorization. To alleviate this problem, we calculate the gradient scale based on the magnitude change of the intermediate image [37]:
$$
s _ {n o r m} = \frac {\left\| x _ {t} - x _ {t - 1} ^ {\prime} \right\| _ {2}}{\left\| g \right\| _ {2}} \cdot s, \tag {8}
$$
where  $x_{t - 1}^{\prime}\sim \mathcal{N}(\mu_{\theta},\Sigma_{\theta})$  . In this way, the dynamic guidance weight  $s_{norm}$  varies along iterations, more effectively guiding the output towards the target, thus improving the output quality.
Second, the conventional classifier guidance typically executes a single gradient step at each denoising step. However, a single gradient step may not sufficiently steer the output toward the intended target, particularly when the intermediate outputs are laden with noise in the early phases of the denoising process. To address this, we allow multiple gradient steps at each denoising step [5] to improve flexibility. Specifically, one can improve the guidance strength of a specific property by increasing the number of gradient steps. The process degenerates to the conventional classifier guidance when the number of gradient steps is set to 1. During inference, users have the flexibility to modulate the strength of guidance for each property as per their requirements, thus boosting overall controllability.
Composite Guidance. Our partial guidance controls only the properties of high-quality outputs, and therefore can be easily extended to complex degradations by stacking respective properties. This is achieved by compositing the classifiers and summing the loss corresponding to each property. An example of composite tasks is shown in Table 1. In addition, we also demonstrate that additional losses such as perceptual loss [2, 16] and adversarial loss [22] can be incorporated for further quality improvement. Experiments demonstrate that our PGDiff achieves better performance than existing works in complex tasks, where accurate modeling of the degradation process is impossible.
CNN/Transformer-based Methods
Diffusion-prior-based Methods
Table 2: Quantitative comparison on the real-world LFW-Test, WebPhoto-Test, and WIDER-Test datasets. Red and blue indicate the best and the second best performance, respectively.
Dataset	Metric	CNN/Transformer-based Methods	Diffusion-prior-based Methods
GFP-GAN [40]	CodeFormer [48]	DiffFace [45]	Ours
LFW-Test	FID↓	72.45	74.10	67.98	71.62
NIQE↓	3.90	4.52	5.47	4.15
WebPhoto-Test	FID↓	91.43	86.19	90.58	86.18
NIQE↓	4.13	4.65	4.48	4.34
WIDER-Test	FID↓	40.93	40.26	38.54	39.17
NIQE↓	3.77	4.12	4.44	3.93
Table 3: Quantitative comparison on the synthetic CelebRef-HQ dataset. Red and blue indicate the best and the second best performance, respectively.
Metric	CNN/Transformer-based Methods	Diffusion-prior-based Methods
GFP-GAN [40]	CodeFormer [48]	DifFace [45]	Ours (w/o ref)	Ours (w/ ref)
FID↓	186.88	129.17	123.18	119.98	121.25
MUSIQ↑	63.33	69.62	60.98	67.26	64.67
LPIPS↓	0.49	0.36	0.35	0.34	0.35
IDS↑	0.36	0.55	0.56	0.44	0.76
By exploiting the diffusion prior, our PGDiff applies to a wide range of restoration tasks by selecting appropriate guidance. In this section, we will introduce the guidance formulation and provide experimental results.
4.1 Blind Face Restoration
Partial Guidance Formulation. The objective of blind face restoration is to reconstruct a high-quality face image given a low-quality input corrupted by unknown degradations. In this task, the most straightforward approach is to train a network with the MSE loss using synthetic pairs. However, while these methods are able to remove the degradations in the input, it is well-known [26] that the MSE loss alone results in over-smoothed outputs. Therefore, extensive efforts have been devoted to improving the perceptual quality, such as incorporating addition losses (e.g., GAN loss) [22, 10, 16, 46, 8] and components (e.g., codebook [48, 13, 43, 47, 35] and dictionary [23, 24, 12, 8]). These approaches often require multi-stage training and experience training instability.
In our framework, we decompose a high-quality face image into smooth semantics and high-frequency details, and provide guidance solely on the smooth semantics. In this way, the output  $\hat{x}_0$  in each diffusion step is guided towards a degradation-free solution space, and the diffusion prior is responsible for detail synthesis. Given an input low-quality image  $y_0$ , we adopt a pre-trained face restoration model  $f$  to predict smooth semantics as partial guidance. Our approach alleviates the training pressure of the previous models by optimizing model  $f$  solely with the MSE loss. This is because our goal is to obtain smooth semantics without hallucinating unnecessary high-frequency details. Nevertheless, one can also provide guidance of various forms by selecting different restorers, such as CodeFormer [48]. The loss for classifier guidance is computed as:  $\mathcal{L}_{res} = ||\hat{x}_0 - f(y_0)||_2^2$ .
Qualitative Results. We evaluate the proposed PGDiff on three real-world datasets, namely LFW-Test [40], WebPhoto-Test [40], and WIDER-Test [48]. We compare our method with both task-specific CNN/Transformer-based restoration models [48, 26, 40] and diffusion-prior-based models[9, 42, 45]. As shown in Fig. 2, existing diffusion-prior-based methods such as GDP [9] and DDNM [42] are unable to generalize to real-world degradations, producing outputs with notable artifacts. In contrast, our PGDiff successfully removes the degradations and restores the facial details invisible in the input
2Among them, GDP [9] and DDNM [42] support only  $4\times$  fixed-kernel downsampling, while DifFace [45] is a task-specific model for blind face restoration.
images. Moreover, our PGDiff performs favorably over task-specific methods even without extensive training on this task.
Quantitative Results on Real-world Datasets. To compare our performance with other methods quantitatively on real-world datasets, we adopt FID [14] and NIQE [27] as the evaluation metrics and test on three real-world datasets: LFW-Test [40], WebPhoto-Test [40], and WIDER-Test [48]. LFW-Test consists of the first image from each person whose name starts with A in the LFW dataset [40], which are 431 images in total. WebPhoto-Test is a dataset comprising 407 images with medium degradations collected from the Internet. WIDER-Test contains 970 severely degraded images from the WIDER Face dataset [40]. As shown in Table 2, our method achieves the best or second-best scores across all three datasets for both metrics. Although GFP-GAN achieves the best NIQE scores across datasets, notable artifacts can be observed, as shown in Fig. 2. Meanwhile, our method shows exceptional robustness and produces visually pleasing outputs without artifacts.
Quantitative Results on Synthetic Dataset. We present a quantitative evaluation on the synthetic CelebRef-HQ dataset [24] in Table 3. Considering the importance of identity-preserving in blind face restoration, we introduce reference-based restoration in Sec. 4.5 in addition to the general restoration in Sec. 4.1. Table 3 shows that our methods achieve best or second best scores across both no-reference (NR) metrics for image quality (i.e., FID and MUSIQ) and full-reference (FR) metrics for identity preservation (i.e., LPIPS and IDS). Since we employ heavy degradation settings when synthesizing CelebRef-HQ, it is noteworthy that identity features are largely distorted in severely corrupted input images. Thus, it is almost impossible to predict an identity-preserving face without any additional identity information. Nevertheless, with our reference-based restoration, we observe that a high-quality reference image of the same person helps generate personal characteristics that are highly similar to the ground truth. The large enhancement of identity preservation is also indicated in Table 3, where our reference-based method achieves the highest IDS, increasing by 0.32.
4.2 Face Colorization