Latest diffusion-based methods for many image restoration tasks outperform traditional models, but they encounter the long-time inference problem. To tackle it, this paper proposes a Wavelet-Based Diffusion Model (WaveDM). WaveDM learns the distribution of clean images in the wavelet domain conditioned on the wavelet spectrum of degraded images after wavelet transform, which is more time-saving in each step of sampling than modeling in the spatial domain. To ensure restoration performance, a unique training strategy is proposed where the low-frequency and high-frequency spectrums are learned using distinct modules. In addition, an Efficient Conditional Sampling (ECS) strategy is developed from experiments, which reduces the number of total sampling steps to around 5. Evaluations on twelve benchmark datasets including image raindrop removal, rain steaks removal, dehazing, defocus deblurring, demoir\'eing, and denoising demonstrate that WaveDM achieves state-of-the-art performance with the efficiency that is comparable to traditional one-pass methods and over 100$\times$ faster than existing image restoration methods using vanilla diffusion models.
翻译:最新的基于扩散模型的方法在许多图像恢复任务中表现优于传统模型,但面临长时推理问题。为此,本文提出了一种基于小波的扩散模型(WaveDM)。WaveDM在小波域中学习干净图像的分布,并基于退化图像经小波变换后的小波谱进行条件建模,这使得每一步采样的时间比在空间域建模更节省。为确保恢复性能,提出了一种独特的训练策略,其中低频和高频谱分别通过不同模块学习。此外,从实验中发展出一种高效条件采样(ECS)策略,将总采样步数减少至约5步。在包括图像雨滴去除、雨条纹去除、去雾、散焦去模糊、去摩尔纹和去噪等十二个基准数据集上的评估表明,WaveDM实现了最先进的性能,其效率可与传统单次方法相比,且比使用原始扩散模型的现有图像恢复方法快100倍以上。