Efficient Image Restoration through Low-Rank Adaptation and Stable Diffusion XL (2024)

2.1 Image Restoration

The purpose of image restoration is to convert degraded images into clean, high-quality images [8, 14, 45, 44]. Typical image restoration problems include super-resolution [23, 6, 17], deblurring [26, 46, 4], and denoising [15, 50, 38], but these methods generally have limited generalization ability, making it difficult to handle degraded images in the real world.

Deep learning introduces other architectures and training paradigms, thereby improving image restoration performance. For example, transformer-based models enhance the fidelity and authenticity of restored images [13]. In addition, attention mechanisms [33] and multi-scale feature extraction techniques are integrated into the restoration framework to better capture fine details of images.

In addition to these methods, diffusion models also receive attention in image restoration tasks. Models such as Denoising Diffusion Implicit Model (DDIM) [32] and Denoising Diffusion Neural Network (DDNM) [36] iteratively refine images through a series of denoising steps, effectively handling various types of image degradation, and show promising results. These models utilize the power of diffusion processes to gradually improve image quality, making them highly effective in tasks such as denoising, deblurring, and more. However, achieving robust performance in various invisible degradation scenarios still faces challenges. Moreover, to better optimize the image restoration model, researchers propose new loss functions, such as Perceptual Loss, which better capture the details of the image and improve the restoration effect. Future research aims to develop more adaptable and scalable models that can effectively generalize to various types of image degradation encountered in real-world applications. Over time, many models that can handle multiple degradation scenarios emerge, among which two-stage methods show good results, such as DiffBIR [19] and SUPIR [42].

2.2 Low Rank Adaptation

Low Rank Adaptation (LoRA) [12] is an approximate numerical decomposition technique that is particularly useful for large-scale language models. This method involves performing low rank decomposition of feature matrices, which allows for efficient adaptation of pre-trained models. By utilizing low rank decomposition techniques, LoRA can significantly reduce the number of parameters in the feature matrix, leading to decreased memory usage and computational overhead.

The core idea behind LoRA is to insert a low rank adaptation matrix into the model architecture, enabling fast adaptation and efficient fine-tuning without altering the original model’s weights. This approach is particularly advantageous in scenarios where computational resources are limited or where rapid model updates are necessary. LoRA achieves this by leveraging the inherent low-rank structure within the feature matrices of large-scale models, which often contain redundant information. By decomposing these matrices into lower-dimensional components, LoRA reduces the complexity of the adaptation process. Furthermore, LoRA’s ability to maintain the integrity of the original model’s weights ensures that the foundational knowledge embedded within the pre-trained model is preserved. This makes LoRA especially effective in transfer learning scenarios, where the pre-trained model is adapted to new tasks or domains. The low-rank adaptation matrix can be trained with relatively fewer resources compared to re-training the entire model, thus making the process more efficient and cost-effective.

In practical terms, LoRA can be applied to various aspects of model adaptation, including fine-tuning for specific tasks, domain adaptation, and even continual learning. Its flexibility and efficiency make it a valuable tool in the toolkit of machine learning practitioners, particularly when dealing with large and complex models. The reduced computational burden also facilitates experimentation and iteration, allowing researchers and engineers to explore a wider range of configurations and settings. LoRA can offer a blend of efficiency, effectiveness, and practicality. Its application can lead to more responsive and adaptable AI systems, capable of quickly incorporating new information and tasks without the need for extensive computational resources.

2.3 Stable Diffusion XL

Diffusion models [11, 27, 40] garner significant attention in the field of generative artificial intelligence, delivering state-of-the-art outcomes across various applications, including text-to-image [16, 29, 48] and text-to-video [3, 43, 37] transformations. These models operate by gradually transforming a simple, structured noise distribution into a complex data distribution through a series of iterative refinement steps. This process enables the generation of high-fidelity images and videos from random noise, making diffusion models a powerful tool for various generative tasks.

Stable diffusion [27, 5, 22] is particularly influential in text-to-image synthesis, leveraging the Latent Diffusion Model (LDM) to execute diffusion operations within a semantically compressed space. This approach enhances computational efficiency by reducing the dimensionality of the data on which diffusion operations are performed. The core architecture of stable diffusion models centers around U-Net, a convolutional neural network architecture that is well-suited for image restoration tasks. U-Net iteratively denoises random latent codes, supported by text encoders and image decoders, to harmonize text and image generation. The use of text encoders allows the model to understand and incorporate textual descriptions into the generated images, resulting in highly detailed and contextually relevant outputs. However, the computational demands of its multi-step inference process [39] become a significant burden, particularly when generating high-resolution images or long video sequences. Each step in the diffusion process requires complex computations, leading to substantial time and resource consumption. This computational overhead poses a challenge for real-time applications and large-scale deployments, where efficiency is crucial. To address these challenges, researchers have introduced various distillation techniques such as Progressive Distillation and Adversarial Distillation [30, 31]. Progressive Distillation incrementally transfers knowledge from a complex model to a simpler one, maintaining the performance while reducing the number of necessary computation steps. Adversarial Distillation, on the other hand, leverages adversarial training to enhance the quality of the distilled model, ensuring that the simplified model retains the generative capabilities of the original.

SDXL Lightning [18] is an enhanced version of the SDXL model that employs progressive adversarial distillation technology, to significantly boost the quality and efficiency of image generation. SDXL Lightning employs an advanced model architecture and adversarial training mechanism to generate high-resolution, detailed images while minimizing computational resources.In order to reduce the computational requirements of training diffusion models for high-resolution image synthesis, it has been found that although diffusion models can ignore perceptually unimportant details through undersampling loss terms, they still require expensive function evaluation in pixel space, thereby requiring a large amount of computation time and energy resources. Therefore, a method was introduced to clearly separate the compression learning phase. This method employs an autoencoder model that learns a space that is perceptually equivalent to the image space, but with significantly reduced computational complexity.

3.1 Background of Stable Diffusion

The key steps of the stable diffusion model consist of the forward diffusion process and the backward denoising process [11]. The forward diffusion process progressively adds noise to the data, whereas the backward process utilizes the learned model to remove noise and restore the original data. Specifically, the forward diffusion process can be described as:

Among them, $\mu_{\theta}$ and $\sigma_{\theta}$are the learned mean and variance functions, respectively.Its goal is to predict a noise to be added to the input image $x$ based on the noisy image $x_{t}$ at time $t$. The objective function of LDM [27] is

In LDM, we learn in latent space, that is, predict the noise added on $\mathbf{z}_{t}$, and the corresponding loss function is expressed as follows:

Compared with traditional Generative Adversarial Networks (GANs) [9], stable diffusion models exhibit enhanced stability and a reduced risk of mode collapse. Moreover, by functioning within the latent space, it substantially enhance the computational efficiency of the generation process.

3.2 Scaling-UP Image Restoration

SUPIR combines large-scale pre trained generative models, significantly improving the effectiveness of image restoration. SUPIR adopts a two-stage architecture, with each stage optimized for different tasks. In the first stage, a pre-trained restoration module is employed to remove degraded components from the image, such as blur and noise. In the second stage, SUPIR leverages sdxl for image detail and texture reconstruction.

First, input a low-quality image, and then the low-quality image will be encoded by the fine tuned encoder and mapped to the latent space. This encoder has been specially trained to handle degraded images. The author designed an adapter based on ControlNet [47] that can recognize LQ image content and guide it to recover images based on the provided low-quality input. The adapter adopts a partially trimmed Vision Transformer (ViT) [1] module and introduces a ZeroSFT [35] module to enhance the guidance effect of LQ images.

In this study, we explore not only the use of SDXL but also several other variants of the Stable Diffusion model. Among these, the SDXL model demonstrates superior performance, while the SDXL Lightning variant also exhibits commendable capabilities. SDXL Lightning is particularly noteworthy for its ability to reduce the number of inference steps to 15 without compromising image quality, thereby enabling the generation of high-quality images within a significantly shorter time frame.

3.3 Training for Low Rank Adaptation

In this study, we adopted LoRA technology to adapt the SDXL model and enhance its performance in facial image generation. LoRA effectively fine tunes the pre trained model by introducing a low rank factorization matrix without increasing memory and computational overhead. Specifically, we incorporated a LoRA adaptation layer into the SDXL model and trained two separate LoRa models. One LoRa model was trained on 1300 landscape images, while the second was specifically trained on 300 facial images. All images are preprocessed to a resolution of 512x512 to ensure data consistency and quality. Each image is labeled in detail to ensure learning ability during training. This method has been proven to be efficient and adaptable under resource constrained conditions in multiple studies [2]. During the training process, the LoRA adaptation matrix is initialized to zero and the update matrix is initialized using a random Gaussian distribution to ensure the stability of the model in the early stages of training. We adjust hyperparameters such as learning rate based on specific performance to optimize model performance. Through this low rank adaptation method, we successfully improved the performance of the SUPIR model in facial image generation and validated the effectiveness of LoRA technology in model fine-tuning.

As mentioned earlier, LoRA represents parameter updates through low rank matrix decomposition. Given a weight matrix $W$, LoRA decomposes it into two low-rank matrices $A$ and $B$, such that:

where the ranks of $A$ and $B$ are much smaller than the original dimensions of $W$.

During training, LoRA updates only the matrices $A$ and $B$, keeping $W$ unchanged. This method drastically reduces the number of parameters that need to be trained. For a weight matrix $W\in\mathbb{R}^{d\times k}$, if we choose a rank $r\ll\min(d,k)$, the number of parameters to be adjusted decreases from $d\times k$ to $(d\times r+r\times k)$. The training process can be described as follows:

4.1 Model Training and Quantitative Comparisons

For training, We use the AdamW optimizer [21] with a learning rate of 0.0001. The training process spans 2 days, with a batch size of 256. In our experiments, the integration of two LoRA modules with the SDXL model in the SUPIR framework demonstrated remarkable improvements in image restoration tasks. We used three SDXL models, namely SDXL, SDXL-lighting, HelloWorld-XL, Among them, HelloWorld XL used 20821 images as the training set, which includes different people and actions, as well as many lifelike animals. Moreover, the image quality of close-up portrait output is better than SDXL. And HelloWorld-XL intentionally includes some low-quality images in the training to enhance the model’s response to negative prompts, which is also why HelloWorld-XL performs well in handling blur and noise. The proposed method achieved excellent performance across multiple metrics, with Peak Signal-to-Noise Ratio (PSNR) values significantly higher than the baseline models, indicating clearer and more accurate restored images. Additionally, the Learned Perceptual Image Patch Similarity (LPIPS) scores were notably lower, reflecting a better perceptual similarity to the ground truth images. Furthermore, the Structural Similarity Index Measure (SSIM) scores were substantially improved, showcasing enhanced structural fidelity and visual quality of the restored images. These outstanding results affirm the effectiveness of our approach in producing high-quality image restorations, making it a promising solution for advanced image restoration applications.

To generate low-quality images to test the performance of our method, we introduced various degradations ranging from simple to complex. For quantitative comparison, we selected the following indicators: complete reference indicators PSNR, SSIM, and LPIPS [49]. Compared with the original SUPIR method, our method has improved in all parameter indicators. Similarly, in Fig.2, it can be seen that our model has indeed achieved good results in face restoration, with some progress compared to SUPIR in certain small details and colors

In terms of details in image restoration, our method demonstrates some better features than the original SUPIR model. For example, in LABEL:fig:first we can see that the texture of the goat’s wool on the trained Lora image is more in line with the texture of the original image. In the image of the little girl, we can see that low-quality images basically do not show the earrings. In the SUPIR model, the earrings are also restored to hair, while the image generated by the trained Lora will show the earrings. Therefore, it can be demonstrated that our method generates high fidelity textures.

4.2 Comparison with Other Methods

We also conducted tests on low-quality images and compared them with other models, such as DiffBIR [19], Stable-SR [34], PASD [41]. We selected the following metrics for quantitative comparison: the reference metrics PSNR, SSIM, LPIPS.In terms of results, our method achieved the best scores on PSNR and SSIM and LPIPS, indicating that our method has higher perceptual similarity between the restored image and the reference image than other methods.

LoRA reduces the complexity of the model through parameter decomposition, thereby reducing time. As shown in Tab.3, the comparison between the original method and our method shows that the LoRA method has improved by nearly 7 seconds compared to before. Compared with the other two models, our method still has the shortest time. However, StableSR requires 200 steps to generate a perfect image and consumes a lot of time. This efficiency gain demonstrates the effectiveness of our approach in handling large-scale models. Moreover, the reduction in computational time does not compromise the quality of the generated images, as evidenced by the consistent performance metrics. Moreover, Fig.3, we can clearly see the differences between the stable SR model and other models, but its performance is not very good. The PASD model performs well in restoring details, such as in case 1. However, PASD has a low ability to restore images with high noise and blur. In case 2, it was unable to restore the windows of distant high-rise buildings and still had noise points in the restored images. In case 3, the restoration of the clock changed its original color.