超高清图像生成新SOTA！清华唐杰教授团队提出Inf-DiT：生成4096图像比UNet节省5倍内存。

清华大学唐杰教授团队最近在生成超高清图像方面的新工作：Inf-DiT，通过提出一种单向块注意力机制，能够在推理过程中自适应调整内存开销并处理全局依赖关系。基于此模块，该模型采用了 DiT 结构进行上采样，并开发了一种能够上采样各种形状和分辨率的无限超分辨率模型。与常用的 UNet 结构相比，Inf-DiT 在生成 4096×4096 图像时可以节省超过 5 倍的内存。该模型在机器和人类评估中均实现了生成超高清图像的SOTA。

选取基于SDXL的超高分辨率上采样Inf-Dit, Dall-e 3和真实图像示例.

相关链接

项目网址:https://github.com/THUDM/Inf-DiT

论文地址：https://arxiv.org/abs/2405.04312

论文阅读

Inf-Dit:对任意分辨率图像进行上采样的内存高效扩散变压器

摘要

近年来，扩散模型在图像生成方面表现出色。然而，由于生成超高分辨率图像（例如 4096 × 4096）时内存的二次方增加，生成的图像的分辨率通常限制为 1024 × 1024。

在这项工作中，我们提出了一种单向块注意机制，可以在推理过程中自适应地调整内存开销并处理全局依赖关系。在此模块的基础上，我们采用 DiT 结构进行上采样，并开发了一个能够对各种形状和分辨率的图像进行上采样的无限超分辨率模型。

综合实验表明，我们的模型在机器和人工评估中都能在生成超高分辨率图像方面实现 SOTA 性能。与常用的 UNet 结构相比，我们的模型在生成 4096 × 4096 图像时可以节省 5 倍以上的内存。

方法

(左)Inf-DiT的总体架构。(右)Inf-DiT的内部结构 块。为了简单起见，我们没有描述最初存在于DiT中的Layernorm。

左图：单向块注意力机制。在我们的实现中，每个块直接依赖于每一层中的三个块：左上角、左侧和顶部的块。右图：Inf-DiT 的推理过程。Inf-DiT 每次根据内存大小生成 n × n 个块。在此过程中，只有后续块所依赖的块的 KV 缓存存储在内存中。

效果

2048×2048分辨率下不同方法的详细定性比较。

4096×4096分辨率下不同方法的详细定性比较。

人的评价结果。参与者对不同的图片进行排序 我们将4到1的分数依次分配给模型，最后计算出 所有结果的平均值。Inf-Dit在所有三个类别中都获得了最高分。

生成迭代上采样的样本。上图:Inf-Dit可以上样 图像本身生成几次，并生成不同频率的细节以相应的分辨率。下图:在128x128点未能产生pupul后在分辨率方面，后续的上采样阶段很难纠正这个错误。

结论

在这项工作中，我们观察到生成超高分辨率图像的主要障碍是模型隐藏状态占用的大量内存。 基于此，我们提出了单向块注意机制（UniBA），它可以通过在块之间执行批量生成来降低空间复杂度。借助UniBA，我们训练了Inf-DiT，这是一种4倍内存效率的图像上采样器，它在生成和超分辨率任务中均实现了最先进的性能。