基于扩散模型的通用图像修复选择性沙漏映射

通用图像修复是一项面向真实世界应用的实用且具潜力的计算机视觉任务。该任务的主要挑战在于同时处理不同的退化分布。现有方法主要利用特定任务的条件（例如提示）来指导模型分别学习不同的分布，这被称为多方映射。然而，由于忽略了不同任务之间的共享信息，该方法不适用于通用模型学习。在本工作中，我们提出了一种基于扩散模型的先进选择性沙漏映射策略，称为 DiffUIR。两项新颖的考量使得我们的 DiffUIR 并非平凡之作。首先，我们为模型配备了强条件指导，以获得扩散模型准确的生成方向（选择性）。更重要的是，DiffUIR 将一个灵活的共享分布项（SDT）优雅且自然地集成到扩散算法中，从而逐渐将不同的分布映射到一个共享分布。在逆过程中，结合 SDT 和强条件指导，DiffUIR 迭代地将共享分布引导至具有高图像质量的特定任务分布（沙漏）。无需繁杂的技巧，仅通过修改映射策略，我们就在五个图像修复任务、22 个基准测试上，在通用设置和零样本泛化设置中取得了 SOTA 性能。令人惊讶的是，仅使用一个轻量级模型（仅 0.89M 参数），我们就能实现出色的性能。源代码和预训练模型可在 https://github.com/iSEE-Laboratory/DiffUIR 获取。