Compression-Aware One-Step Diffusion Model for JPEG Artifact Removal

TL;DR

CODiff is a novel one-step diffusion model that effectively removes JPEG artifacts by incorporating compression priors and dual learning strategies, achieving superior results with reduced computational cost.

Contribution

The paper introduces CODiff, a compression-aware one-step diffusion model utilizing a new visual embedder and dual learning to improve JPEG artifact removal.

Findings

Outperforms recent methods in quality metrics

Effectively handles highly compressed JPEG images

Reduces computational overhead of diffusion models

Abstract

Diffusion models have demonstrated remarkable success in image restoration tasks. However, their multi-step denoising process introduces significant computational overhead, limiting their practical deployment. Furthermore, existing methods struggle to effectively remove severe JPEG artifact, especially in highly compressed images. To address these challenges, we propose CODiff, a compression-aware one-step diffusion model for JPEG artifact removal. The core of CODiff is the compression-aware visual embedder (CaVE), which extracts and leverages JPEG compression priors to guide the diffusion model. We propose a dual learning strategy that combines explicit and implicit learning. Specifically, explicit learning enforces a quality prediction objective to differentiate low-quality images with different compression levels. Implicit learning employs a reconstruction objective that enhances the model's generalization. This dual learning allows for a deeper and more comprehensive understanding of JPEG compression. Experimental results demonstrate that CODiff surpasses recent leading methods in both quantitative and visual quality metrics. The code is released at https://github.com/jp-guo/CODiff.