Unleashing the Potential of Transformer Flow for Photorealistic Face Restoration

1 Introduction

Blind face restoration aims to recover high-quality face images from low-quality inputs afflicted by unknown degradations such as low resolution[Zeng et al., 2023], artifacts[Zhang et al., 2022], noise[Yang et al., 2020], and blur[Lai et al., 2022]. Early method[Zhu et al., 2016] directly trained neural networks to regress high-quality face images. In recent years, the quality of restoration has significantly improved through the utilization of various face priors. Geometric priors, including face landmarks[Chen et al., 2018], parsing maps[Chen et al., 2021], and heatmaps[Yu et al., 2018], are essential for accurately restoring the shapes of face components. Generative priors have also played a crucial role;

∗Corresponding author

for instance, methods[Wang et al., 2021a] leveraging StyleGAN learn to restore fine details, while approaches[Zhou et al., 2022a] utilizing VQGAN model priors enhance the overall quality of face restoration.

Currently, diffusion-based face restoration methods have achieved state-of-the-art performance, typically comprising two stages. The first stage employs neural networks such as SwinIR[Liang et al., 2021] to remove artifacts and noise, while the second stage controls the generation of image content based on the low-quality input. For example, DiffBIR[Lin et al., 2023] leverages Stable Diffusion v1-5's conditional generation capabilities to perform face restoration using a U-Net model for latent space denoising. Similarly, FlowIE[Zhu et al., 2024] constructs Noise-GT paired datasets using pre-trained Stable Diffusion models and learns image restoration through a second Rectified Flow[Liu et al., 2022].

Despite these advancements, existing methods exhibit several limitations:

Insufficient texture & lack of realism: Previous methods rely on the U-Net architecture as a score (noise, vector) estimation network. As the number of model parameters increases, the performance of U-Net tends to plateau [Peebles and Xie, 2023a], which makes it difficult to restore more detailed faces. Moreover, the results produced by the Unetbased diffusion face restoration method tend to have an oil painting texture and are not realistic.

Challenges in Training ControlNet: Previous methods trained ControlNet from scratch, which is feasible for models with a small number of parameters. However, as model size continue to grow, training ControlNet from scratch on specific datasets becomes increasingly difficult and costly, making them prone to mode collapse and overfitting. As illustrated in Fig. 2, simply training a ControlNet from scratch for face restoration leads to mode collapse issues.

Suboptimal Training Strategies: The loss function weighting strategies are directly inherited from the original diffusion model[Esser et al., 2024]. However, face restoration task require not only the generation of fine details but also the balance of aesthetic quality. Therefore, directly adopting existing training strategies is ineffective for training gradient estimators tailored to this specific task.

To address the challenges in face restoration, we propose OmniFace, a novel framework built around three core innovations: (1) utilizing the scaling law properties of Transformer with large parameter counts to enhance face texture quality, (2) Propose C-Projector to control the Transformer diffusion model for face restoration, and (3) designing adaptive loss constraints specifically tailored to face restoration task.

Enhancing Face Texture with Transformer Scaling Law The first innovation in OmniFace lies in leveraging the scaling laws of Transformer, which enable significant improvements in texture quality for high-resolution face restoration. Compared to GAN-based and diffusion-based methods, OmniFace fully exploits the large parameter counts of Transformer, achieving superior clarity and realism, as shown in Fig. 1. OmniFace takes full advantage of the scaling laws, delivering high-fidelity face textures.

C-Projector: Effective Control via Parameter Reuse- Based ControlNet The second innovation addresses a key challenge in controlling Diffusion Transformer (DIT) for face restoration. As shown in Fig. 2, direct control of DIT often leads to mode collapse and inconsistent results, limiting their applicability to face restoration task. To overcome this, we propose a parameter reuse-based ControlNet, which effectively stabilizes and controls the denoising Transformer. By reusing parameters, the ControlNet ensures robust operation and prevents mode collapse, enabling the generation of consistent and realistic face detail.

Adaptive Training Strategy We propose an adaptive training strategy that dynamically balances perceptual constraints and structural constraints. Different from the original loss function weighting strategy in the diffusion model, we count the values of each denoising step in the diffusion process on multiple loss functions, and adjust the corresponding loss weights according to the magnitude of the loss function at different time steps, so as to better constrain the trajectory of the face diffusion flow.

OmniFace unifies these innovations into a cohesive framework for face restoration. By fully leveraging Transformer scaling laws, implementing the parameter reuse-based ControlNet, and adopting adaptive vector field constraints, OmniFace achieves state-of-the-art face restoration performance.

2 RelatedWorks

2.1 Deep Learning for Image Enhancement

Deep learning has shown remarkable progress in image enhancement tasks, including inpainting, super-resolution, HDR reconstruction, deraining, and artifact removal. Various methods have been proposed using CNNs, Transformers, and GNNs to improve visual quality. Lightweight and edge-guided models have been developed for image inpainting [Li et al., 2020], while temporal and attention-based techniques have advanced HDR video generation [He et al., 2022; Xu et al., 2023; Zhang et al., 2024b; Zhang et al., 2024a; Zhang et al., 2023]. Graph-based methods explore patch similarities for effective deraining [Wang et al., 2024]. Other works focus on video enhancement, compression, and debanding using adaptive and feature-aware networks [He et al., 2021; Xu et al., 2024b; Xu et al., 2024a].

2.2 GAN-basedMethods

Generative Adversarial Networks (GANs) have been widely explored for providing high-quality face priors in face restoration task. Methods such as [Zhou et al., 2022b; Wang et al., 2021b; Gu et al., 2022] effectively utilize GAN-based priors to address blind face restoration (BFR) and achieve satisfactory results. However, GAN-based approaches often suffer from instability during training and require meticulous hyperparameter tuning, limiting their practical applicability.

2.3 Diffusion-basedMethods

Diffusion models have emerged as a powerful generative framework due to their stability and ability to produce highquality, diverse images [Ho et al., 2020]. Unlike GANs, they avoid mode collapse by iteratively denoising random noise. Recent works have also improved their sampling efficiency without compromising visual fidelity.

In face restoration, diffusion models have shown strong performance in both zero-shot [Kawar et al., 2022] and supervised settings [Lin et al., 2023]. Methods like FlowIE further enhance inference speed by distilling large diffusion models through flow matching strategies.

However, existing diffusion-based face restoration methods often lack the scalability and detail-preserving capabilities of convolutional architectures. To address these issues, we propose an adaptive training strategy that leverages the scaling benefits of Transformers, achieving more realistic and detailed face restoration results.

3 Method

In this section, we introduce OmniFace, a face enhancement method based on Transformer Flow Matching. This method proposes a diffusion representation reuse ControlNet C-Projector to effectively control Transformer to restore high-quality faces. We begin by providing a brief overview of the background on Transformer Diffusion and Flow Matching, followed by a detailed description of our designed conditional network, adaptive training strategy. The overall network architecture of the proposed OmniFace method is illustrated in the Fig.3.

3.1 Preliminary

Diffusion Transformer

In [Peebles and Xie, 2023b], the authors propose Diffusion Transformer (DIT), which integrates transformer architectures into the diffusion modeling framework to enhance scalability and performance in generative task. By leveraging the strengths of transformers, such as modeling long-range dependencies and capturing complex data distributions, DIT addresses the limitations of traditional CNN-based diffusion models.

The core of DIT is a Transformer-Based Denoising Network, which replaces the conventional CNN denoiser with a transformer architecture that utilizes multi-head self-attention to capture global context and intricate patterns. Positional Encoding Integration ensures the retention of spatial information, crucial for task like face restoration [Dosovitskiy, 2020]. Additionally, its Scalable Architecture Design enables efficient handling of larger models and higher-resolution inputs with minimal computational overhead, making DIT a powerful and scalable framework for generative modeling.

Flow Matching

Flow Matching models are generative models that use Normalizing Flows (NFs) to transform complex distributions into simple ones via invertible transformations. Continuous Normalizing Flows (CNFs) extend this idea by modeling these transformations with ordinary differential equations (ODEs), making it easier to capture time-varying data distributions.

In Flow Matching (FM), the goal is to learn a vector field v(x, t) that describes the evolution of data points over time. The data generation process is represented as:

where x(t) is the data at time t, and v(x(t), t) is the vector field governing its evolution. The training objective minimizes the difference between the real data trajectory and the model's trajectory, often using a path-matching loss:

where ˆ v(x(t), t) is the vector field from real data. After training, new samples can be generated by solving the ODE:

This method improves the stability and efficiency of training, particularly in applications like diffusion models, where it helps refine the data generation process.

3.2 Framework

We propose a novel face restoration network, illustrated in Fig. 3. The architecture consists of two main components: a Diffusion Transformer Network and a projection-based ControlNet, called the C-Projector.

To leverage the Transformer's scaling laws for enhanced texture generation, we incorporate a large-scale Denoising Transformer (DIT) into our framework. The DIT's strong detail generation capabilities in high-parameter settings enable the restoration of intricate face features, ensuring high quality face.

To effectively control the Transformer architecture, we design the C-Projector Fc. This ControlNet project the multimodal hidden features hl into control features c, allowing fine-grained manipulation of the restoration process. By utilizing parameter reuse strategies, the C-Projector can efficient training and rapid convergence.

Let the input low-quality image be denoted as x ∈ RH×W×C . The image x is first processed by the prerestoration network Fp, resulting in an initial restoration feature:

Subsequently, ˜x is encoded by the encoder E of a Variational Autoencoder (VAE) to obtain the latent feature zl:

The latent feature zl, random noise ϵ ∈ Rh×w×c , text condition p ∈ Rm, and time step t ∈ R are sequentially input into N multimodal Transformer modules {Ti}Ni=1 , generating hidden features h(i) l :

where h0l = zl. Each hidden feature hil is then projected by the Projector module Fc to obtain the control feature ci :

Within the Denoising Transformer model DT , noise ϵ, text condition p, and time step t are input into N Transformer layers. The output of each layer yi is combined with the control feature c and input to the next Transformer layer:

where y0 = ϵ. The final output yn represents high-quality latent features, which are subsequently decoded by the VAE decoder D to produce the restored image:

This architecture combines the powerful generative capabilities of Transformer model with the efficient control mechanism provided by C-Projector, achieving high-fidelity face image restoration.

3.3 C-Projector: Efficient Control via Parameter Reuse

During the training of Transformer-based conditional diffusion model on the FFHQ dataset, we noticed significant texture anomalies. The original models tended to generate excessive textures, which degraded restoration quality. Specifically, the hidden features of trained model exhibited large weights, a sign of overfitting (Fig. 2).

To solve this, we introduce the C-Projector Fc, a ControlNet that suppresses texture anomalies. The C-Projector uses a parameter reuse strategy by initializing the weights of its N Transformer modules with those from a pre-trained Denoising Transformer T .

Each Transformer module in the C-Projector add a Projector layer Pi that maps low-quality image features x to guidance features ci in the diffusion space:

whereWi and bi are learnable parameters.

Additionally, we apply Low-Rank Adaptation (LoRA) to the Transformer modules in Fc to further improve control with a lightweight model. LoRA introduces low-rank matrices Ai and Bi to update the model weights as follows:

The LoRA method can further enhance the details and realism of the restored face.

3.4 Time Adaptive Loss Weighting Strategy for Face Restoration in Diffusion Model

We propose an adaptive loss weighting strategy for face restoration task in diffusion models. Traditional diffusion models use fixed loss weights, which are not ideal for task requiring specific attributes, such as face identity and finegrained details. To address this, we introduce a dynamic mechanism that adjusts the loss weights based on the model's performance at each timestep, ensuring that critical aspects of face restoration, like perceptual quality and structural accuracy, are prioritized.

Loss Function and Dynamic Weight Adjustment

Traditional diffusion models rely on static loss weights (e.g., MSE) across all timesteps. However, these fixed weights are not well-suited for face restoration, which requires a more flexible approach. Our method incorporates four key loss functions: mean squared error (MSE), face recognition loss, perceptual loss (LPIPS), and structural similarity loss (SSIM). The weight for each loss function is adaptively adjusted based on its magnitude at each timestep.

Formally, at timestep t, the total loss is:

where ˆx, x respectively predict the latent and high-quality latent, and the weights are computed as:

where E[Lloss] is the average of the respective loss function across previous timesteps. This ensures that the loss weights are adapted based on the historical performance of each loss term.

Time-Step Dependent Weighting

Different stages of the restoration process require different focuses: early timesteps prioritize structural restoration, while later steps refine identity and details. Our adaptive strategy adjusts the weights based on the loss magnitudes, allowing the model to prioritize important features as needed.

For example, when face recognition loss is lagre, the loss function increases the weight on identity preservation, ensuring that the face's identity is preserved as the restoration progresses. Similarly, if perceptual or structural losses are significant, the model adapts to refine the generated texture and structure. In the case of MSE, when large pixel discrepancies exist, the model increases the MSE weight to focus on overall pixel-level accuracy.

Advantages of Adaptive Weighting

The adaptive loss weighting strategy allows the model to dynamically focus on the most critical aspects of face restoration, improving both the efficiency and quality of the restoration process. This approach facilitates faster convergence and enhances identity preservation and face detail recovery. Ultimately, our method enables more precise and high-quality face restoration in diffusion models, making it particularly effective for task that require both structural accuracy and finegrained preservation.

4 Experiments

4.1 Datasets, Metrics

Train Dataset. We train our model using the FFHQ dataset[Karras, 2019], which contains 70,000 high-quality (HQ) face images. During training, all images are resized to 1024×1024. To generate low-quality (LQ) images for training pairs, we apply a degradation model. Specifically, we first convolve the HQ image with a Gaussian kernel to introduce blur, followed by downsampling by a factor of r. Gaussian noise with strength δ is then added, and JPEG compression with quality factor q is applied. Finally, the LQ image is resized back to 1024×1024. The parameters for the degradation process, including the Gaussian kernel size σ, downsampling ratio r, noise strength δ, and JPEG quality factor q, are randomly sampled from predefined ranges to ensure diversity and complexity in the training data. The degradation model is consistent with [Zhou et al., 2022a].

Test Datasets. Due to the lack of high-resolution face datasets, we use the last 500 images from the FFHQ dataset as test set, which are also resized to 1024×1024. To evaluate the model under various degradation conditions, we apply the same degradation process to these test images as used during training. Additionally, we evaluate our method on several other datasets, including CelebA-HQ[Karras, 2017], LFW-Test[Wang et al., 2021b], WebPhoto-Test[Wang et al.,

2021b], WIDER-Test[Zhou et al., 2022b], and Child[Wang et al., 2021b]. Although the images in these datasets are of lower resolution, they cover varying levels of degradation, ranging from mild to severe, providing a comprehensive evaluation of the model's performance under different real-world conditions.

Evaluation Metrics. For datasets with ground truth (GT) images, we use PSNR, SSIM, and LPIPS as the basic evaluation metrics. These metrics provide a quantitative assessment of image quality, with PSNR measuring pixel-wise reconstruction accuracy, SSIM capturing structural similarity, and LPIPS evaluating perceptual similarity. In addition, to ensure comprehensive evaluation, for datasets without ground truth, we employ several no-reference metrics, including MUSIQ[Ke et al., 2021], MANIQA[Yang et al., 2022], NIQE[Mittal et al., 2013], CLIPIQA[Wang et al., 2023], FID[Jayasumana et al., 2024], NIMA[Talebi and Milanfar, 2018] and TOPIQ[Chen et al., 2024], to assess image quality. These metrics provide insight into the overall quality and naturalness of the generated images without requiring ground truth data. Furthermore, we also evaluate the realistic and aesthetic quality[Wu et al., 2024] of the restored faces generated by different methods. This subjective assessment ensures a more holistic evaluation of the models' performance in terms of generating realistic and visually appealing face images.

4.2 Comparisons with State-of-the-Art Methods

We compared the proposed method with state-of-the-art approaches, including RealESRGAN (GAN-based), CodeFormer (VQGAN-based), and PASD (Diffusion-based). Extensive evaluations were conducted on both synthetic and real-world datasets.

Quantitative Results. Table 1, Table 2 and Table 3 presents quantitative comparison on synthetic and real datasets. Our proposed OmniFace outperforms existing methods in terms of image quality metrics such as FID, and

MUSIQ. Moreover, OmniFace achieves state-of-the-art performance on perceptual evaluation metrics, including CLIPIQA and TOPIQ, demonstrating the superior subjective quality of the generated results. Furthermore, we used the latest TOPIQ(Face) metric to evaluate face quality. The results show that the proposed method significantly outperforms previous approaches in generating high-quality face reconstructions, as shown in Table 6. To better evaluate the quality of results produced by different methods, we introduced VLMbased Q-Align for image quality assessment. The proposed method achieves the best performance on the Q-Align metric, as shown in Table 4.

Visual Results. Fig. 1 provides a qualitative comparison, highlighting the limitations of competing methods. For example, CodeFormer introduces artifacts around the eyes, and PASD generates less realistic skin textures. Additionally, all baseline methods struggle to produce photorealistic hair. Leveraging the detail generation capabilities of Transformer, our proposed method excels in producing highly realistic hair and skin textures, achieving visually compelling results.

FFHQ CelebA

512

Child LFW Web

Photo

RealESRGAN0.23 0.20 0.25 0.47 0.48 0.23 0.08

CodeFormer 0.68 0.66 0.70 0.51 0.67 0.59 0.61

0.79 0.76 0.82 0.66 0.77 0.72 0.68

OnmiFace 0.82 0.79 0.90 0.82 0.87 0.78 0.85

Table 6: Quantitative comparison of methods using TOPIQ(Face), a specialized metric for assessing face quality, across multiple benchmarks.

4.3 Aesthetic evaluation

We further evaluated the aesthetic quality of the results produced by the proposed method using the Q-Align(Aesthetic) metric. As shown in Table 5, the proposed method achieves the best performance across all datasets, demonstrating its ability to generate visually pleasing and high-quality outputs.

4.4 Ablation studies

To evaluate the contributions of the proposed C-Projector ControlNet and the adaptive loss strategy, we conducted extensive ablation experiments. Quantitative results, shown in

Wider Table 7, demonstrate that removing either component leads to a significant drop in performance across multiple evaluation metrics, including MUSIQ, CLIP-IQA, and TOPIQ. Furthermore, qualitative comparisons in Fig. 5 and Fig. 6 highlight the importance of these modules. The removal of the CProjector results in noticeable artifacts, while excluding the adaptive loss strategy compromises the realism and richness of textures, particularly in hair and skin. These findings validate the effectiveness of both components in enhancing face restoration quality.

W C-Projector

Figure 5: Qualitative comparison of face restoration results with and

without the C-Projector.

Figure 6: Comparison of face restoration results without (left) and with (right) the proposed adaptive loss. Our adaptive loss enables the recovery of richer details, such as the hair texture, while maintaining overall consistency and realism.

5 Conclusion

We proposed OmniFace, a Transformer Flow based face restoration framework that leverages diffusion flow to achieve high-resolution, photorealistic results. OmniFace outperforms state-of-the-art methods across multiple benchmarks and metrics, with ablation studies confirming the importance of its C-Projector and adaptive loss strategy. Our work highlights the potential of scaling Transformer for face restoration and sets a new benchmark in realistic face restoration.

Acknowledgements

This paper is supported by the construction project of the Fujiang Laboratory Nuclear Medicine Artificial Intelligence Research Center (No. 2023ZYDF074) , supported by the Fundamental Research Funds for the Central Universities under Grant ZYTS25270. Kepeng Xu and Li Xu contributed equally to this work. Please ask Dr. Gang He for correspondence.