【深度学习】UniControl 一个统一的扩散模型用于可控的野外视觉生成

论文：https://arxiv.org/abs/2305.11147
代码：https://github.com/salesforce/UniControl#data-preparation
docker快速部署：https://qq742971636.blog.csdn.net/article/details/133129146

Abstract

在设计交互式AI系统时，实现机器自主性和人类控制往往代表了不同的目标。像Stable Diffusion这样的视觉生成基础模型在解决这些目标方面显示出潜力，特别是当以任意语言（arbitrary languages）进行提示时。然而，它们通常在生成具有空间、结构或几何控制的图像方面表现不佳。将这些控制集成到一个统一模型中，以适应单一模型中的各种视觉条件，仍然是一个未解决的挑战。因此，我们引入了UniControl，这是一个新的生成基础模型，它将各种可控的条件到图像（C2I）任务集成到一个统一的框架中，同时仍然允许使用任意语言提示。UniControl实现了像素级精确的图像生成，其中视觉条件主要影响生成的结构，而语言提示则指导了风格和上下文（context）。为了使UniControl具备处理各种视觉条件的能力，我们增强了预训练的文本到图像扩散模型，并引入了一个任务感知的HyperNet来调节扩散模型，使其能够同时适应不同的C2I任务。在九个独特的C2I任务上训练后，UniControl展示出了在看不见的视觉条件下的令人印象深刻的零样本生成能力。实验结果显示，UniControl通常超越了具有相似模型大小的单一任务控制方法的性能。这种控制多功能性将UniControl定位为可控视觉生成领域的重要进展。

Introduction

生成基础模型正在彻底改变人与人工智能在自然语言处理（NLP）[1–6]、计算机视觉（CV）[7–10]、音频处理（AP）[11, 12]和机器控制[13–15]等领域的交互方式。在NLP领域，像InstructGPT或GPT-4这样的生成基础模型在广泛的任务上表现出色，例如问题回答、摘要、文本生成或机器翻译，都可以在单一统一模型中实现。这种多任务能力是生成基础模型最吸引人的特征之一。此外，生成基础模型还可以在未见过的任务上进行零样本或少样本学习[3, 16, 17]。

对于视觉领域的生成模型[9, 18–20]，这种多任务能力较不明显。稳定扩散模型（SDM）[9]已经成为文本条件下图像生成的主要基础。然而，虽然文本描述提供了一种非常灵活的方式来控制生成的图像，但它们往往无法提供对空间、结构或几何控制的像素级精确度。最近的一个工作，ControlNet [21]，旨在增强SDM以实现对视觉条件（例如边缘图、深度图）的控制。通过额外的视觉条件，ControlNet可以实现对生成结构的明确空间、结构或几何控制，而不会失去来自文本说明的语义控制。不幸的是，不像一个统一模块如CLIP [22]可以处理的语言提示，每个ControlNet模型只能处理其训练过的特定控制模态（例如边缘图）。需要重新训练一个单独的模型来处理不同模态的视觉条件，这带来了不小的时间和空间复杂性成本。

为了克服以前工作的限制，我们提出了UniControl，这是一种用于野外可控视觉生成的统一扩散模型，能够同时处理语言和各种视觉条件。自然地，UniControl可以进行多任务处理，并可以将来自不同任务的视觉条件编码为通用表示空间，寻求在任务之间找到共同的表示结构。UniControl的统一设计使我们能够享受到改进的训练和推理效率以及增强的可控生成。

一方面，UniControl的模型大小在任务数量增加时并不显著增加。另一方面，UniControl从不同视觉条件之间的固有联系中获得了优势[例如，23–25]。这些关系，如深度和分割映射，利用共享的几何信息来提高可控生成的质量。

UniControl的统一可控生成能力依赖于两个新设计的模块，即混合专家（MOE）风格适配器和任务感知的HyperNet。MOE风格适配器可以从各种视觉条件中学习必要的低级特征图，使UniControl能够捕捉不同视觉条件的独特信息。任务感知的HyperNet以任务指令作为自然语言提示输入，并输出任务感知的嵌入。输出的嵌入可以用于调节ControlNet [21]以实现任务感知的视觉条件控制，其中每个任务对应于特定格式的视觉条件。因此，任务感知的HyperNet允许UniControl学习跨越各种任务的元知识，并具备泛化到未见任务的能力。

为了获得多任务处理和零样本学习的能力，我们在五个类别的九个不同任务上对UniControl进行了预训练：1）边缘（Canny、HED、用户草图）；2）区域映射（分割映射、边界框）；3）骨架（人体姿势骨架）；4）几何映射（深度、表面法线）；5）编辑（图像外延）。我们建立了MultiGen-20M数据集，包括超过2000万高质量的原始图像、语言提示和所有任务的视觉条件的三元组。然后，UniControl在NVIDIA A100硬件上进行了超过5000小时的训练。结果，UniControl展现出了适应新任务的零样本能力，突显了其在实际应用中的潜力。我们的贡献总结如下：

• 我们提出了UniControl，这是一个统一的模型，能够处理各种视觉条件以进行可控的视觉生成。

• 我们收集了一个新的多条件视觉生成数据集，包括来自五个类别的九个不同任务的超过2000万个图像文本条件三元组。

• 我们进行了广泛的实验，以证明统一模型UniControl优于每个单一任务控制的图像生成，这要归功于学习不同视觉条件之间的内在关系。

• UniControl表现出了在零样本方式下适应新任务的能力，突显了其在野外广泛应用的多功能性和潜力。

Related Works

基于扩散的生成模型。扩散模型最初是在[26]中引入的，对生成图像产生了有利的结果[18, 21]。通过各种训练和采样技术，如基于分数的扩散[27, 28]、去噪扩散概率模型（DDPM）[29]和去噪扩散隐式模型（DDIM）[30]，已经进行了改进。在训练高分辨率图像的U-Net去噪器[31]时，研究人员采用了包括金字塔[32]、多个阶段[20]或潜在表示[9]在内的加速技术。特别是，UniControl利用稳定扩散模型（SDM）[9]作为基本模型进行多任务处理。

文本到图像扩散。扩散模型在文本到图像生成任务中取得了尖端性能[20, 19]，通过在扩散生成器中使用CLIP [22]或T5预训练[2]文本嵌入，与扩散生成器中的交叉注意力的U-Net去噪器。GLIDE [33]是另一个支持图像生成和编辑的文本引导扩散模型的例子。UniControl和与之密切相关的ControlNet [21]都是基于以前的基于扩散的文本到图像生成工作[9]构建的。

图像到图像翻译。图像到图像（I2I）翻译任务最初是在Pix2Pix [34]中提出的，重点是学习不同领域中图像之间的映射。最近，基于扩散的方法[35, 36, 21]建立了最新的技术成果。最近的基于扩散的图像编辑方法在不需要配对数据的情况下表现出色，例如SDEdit [37]、prompt-to-prompt [38]、Edict [39]等。其他图像编辑示例包括各种扩散桥和流[40–44]、基于分类器引导[28]的方法用于上色、超分辨率[32]、修复[45]等。ControlNet [21]同时采用了视觉和文本条件，并实现了新的最先进的可控图像生成。我们提出的UniControl统一了ControlNet的各种视觉条件，并能够在新的未见任务上进行零样本学习。与此同时，Prompt Diffusion [46]引入了来自图像修复的视觉提示[47]到可控扩散模型中，这需要两对额外的图像作为训练和推理的上下文示例。相比之下，UniControl仅使用单一视觉条件，同时仍然能够进行多任务处理和零样本学习。

UniControl

Training Setup

与以前的生成模型（如稳定扩散模型（SDM）[9]或ControlNet [21]）不同，这些模型的图像生成条件是单一的语言提示，或者单一类型的视觉条件，如Canny。

UniControl需要处理来自不同任务的各种视觉条件，同时还需要处理语言提示。

Model Design

由于我们的统一模型UniControl需要在一系列不同的任务上实现出色的性能，因此有必要确保网络设计具备以下属性：1）模型能够克服来自不同任务的低级特征的不匹配；2）模型能够学习跨任务的元知识，并有效适应每个任务。

第一个属性可以确保UniControl能够从所有任务中学习到必要且独特的信息。例如，如果UniControl将分割地图作为视觉条件，那么模型可能会忽略3D信息。结果，学到的特征图可能不适用于以深度图像作为视觉条件的任务。第二个属性将允许模型学习跨任务的共享知识，以及它们之间的差异。

我们引入了两个新设计的模块，MOE风格适配器和任务感知的HyperNet，使UniControl能够具备上述两个属性。UniControl的模型设计概述如图2所示。我们将在下面详细描述这些模块的设计。

任务泛化能力

通过在MultiGen-20M数据集上进行全面预训练，UniControl展示出了对在其训练过程中未遇到的任务具有零样本能力，这表明UniControl具备超越域内分布以实现更广泛泛化的能力。我们在以下两种情况下展示了UniControl的零样本能力：

混合任务泛化。如图3左侧所示，我们考虑将两种不同的视觉条件作为UniControl的输入，即分割地图和人体骨架的混合组合，并在文本提示中加入特定的关键词“背景”和“前景”。

此外，我们将混合任务指令重写为合并两个任务的指令的混合，例如“分割地图和人体骨架到图像”。

零样本新任务泛化。如图3右侧所示，UniControl需要在新的未见过的视觉条件上生成可控图像。为实现这一目标，基于未见和已训练任务之间的关系估算任务权重至关重要。任务权重可以通过手动分配或在嵌入空间中计算任务指令的相似性得分来估算。图5（d）中的示例结果是由我们手动分配的MOE权重生成的，例如“深度：0.6，分割：0.3，Canny：0.1”用于上色。MOE风格适配器可以与估算的任务权重线性组合，以从新的未见视觉条件中提取浅层特征。

实验

我们在各种条件和任务下进行了一系列综合实验，利用不同的数据集来挑战模型的适应性和多功能性，以实证评估UniControl的有效性和鲁棒性。实验设置、方法论和结果分析将在接下来的部分提供。

实验设置

实施。UniControl如图2所示，由稳定扩散、ControlNet、MOE适配器和任务感知的HyperNet组成，总共约1.5B个参数。MOE适配器由并行的卷积模块组成，每个模块对应一个任务。任务感知的HyperNet接受任务指令的CLIP文本嵌入[22]并输出任务嵌入以调制零卷积核的权重。我们基于PyTorch Lightning [52]采用AdamW [51]作为优化器。学习速率设置为1×10^-5。我们的完整版本UniControl模型在16个Nvidia-A100 GPU上进行训练，批处理大小为4，需要约5000个GPU小时。我们还应用了结果的安全检查器作为保障措施。

数据收集。由于ControlNet的训练集目前不可用，我们从头开始启动了自己的数据收集过程，并将其命名为MultiGen-20M。我们使用LaionAesthetics-V2 [53]的一个子集，其中的审美评分超过六，排除小于512的低分辨率图像。这产生了约280万个图像文本对。随后，我们对这个数据集进行处理，创建了五个类别的九个不同任务（边缘、区域、骨架、几何图、真实图像）：

• Canny（2.8M）：使用随机阈值的Canny边缘检测器[54]。
• HED（2.8M）：采用Holistically-nested边缘检测[55]进行鲁棒的边界确定。
• Depth（2.8M）：采用Midas[56]进行单目深度估计。

• Normal（2.8M）：使用深度任务的深度估计结果来估计场景或对象的表面法线。

• 分割（2.8M）：使用在ADE20K [58]数据集上预训练的Uniformer [57]模型生成包括150个类别的分割图。

• 物体边界框（874K）：利用在COCO [60]数据集上预训练的YOLO V4 [59]进行80个物体类别的边界框标注。

• 人体骨架（1.3M）：采用预训练的Openpose [61]模型从源图像生成人体骨架标签。

• 图像外扩（2.8M）：为源图像创建边界蒙版，蒙版的百分比从20%到80%随机变化。

随后，我们对HED图使用高斯滤波和二值阈值处理进行进一步处理，以模拟用户的素描。总体而言，我们积累了超过2000万个图像-提示-条件三元组。任务说明自然地从相应的条件中派生，每个任务对应于一个特定的指令，例如"canny边缘到图像"适用于canny任务。我们保持了任务和指令之间的一对一对应关系，没有引入变化，以确保训练过程的稳定性。此外，我们还收集了一个测试数据集，用于评估每个任务的100-300个图像条件提示三元组。源数据来自Laion和COCO。我们将开源我们的训练和测试数据，以贡献给社区。

基准模型。与UniControl的最直接比较来自于特定任务的ControlNet模型。六个任务与ControlNet中呈现的任务重叠，因此我们选择这些任务的官方模型作为这些任务的基线。为了进行公平比较，我们使用我们收集的数据重新实现了ControlNet模型（单任务）。我们的统一多任务UniControl与每个任务的这些任务感知模型进行比较。我们采用默认的采样器作为DDIM [30]，其中包括导向权重9和步数50。为了确保公平性，用于比较的所有单任务模型都经过了100K次迭代的训练，而我们的多任务模型每个任务的训练也在900K左右，以确保公平性。我们提出的模型的高效性和紧凑设计在其构建中是显而易见的。UniControl的总大小约为1.5B个参数，而单一任务的ControlNet+SDM占用1.4B。为了实现相同的九个任务功能，单一任务策略将需要集成一个SDM和九个任务特定的ControlNet模型，总共约为4.3B个参数。

视觉比较

在图4中，我们对不同任务（Canny、HED、Depth、Normal、Segmentation、Openpose、Bounding Box和Outpainting）进行了视觉比较。我们的方法始终优于基线ControlNet模型，无论是在视觉质量还是在与条件或提示的对齐方面。

对于Canny任务，我们的模型生成的结果在细节保留和视觉一致性方面表现出更高的程度。与ControlNet相比，UniControl的输出保持了对边缘信息的忠实再现（即圆桌）。在HED任务中，我们的模型有效地捕捉到了强大的边界，导致视觉上吸引人的图像，具有清晰而锐利的边缘过渡，而ControlNet的结果似乎不是事实。此外，我们的模型对深度图和表面法线的三维几何引导具有更为微妙的理解，比ControlNet产生了明显更准确的输出。在Normal任务中，我们的模型忠实地再现了法线表面信息（即滑雪杆），从而产生了更加逼真和视觉上优越的输出。在分割、Openpose和边界框任务中，我们模型生成的图像与给定的条件更好地对齐，确保对输入提示的更高保真度。例如，重新实现的ControlNet-BBox误解了“雕像旁边的女人”，而我们的输出显示出高度的准确性和细节。在Outpainting任务中，我们的模型通过生成具有平滑过渡和自然纹理的合理图像来展示其优越性。它胜过了ControlNet模型，后者产生了较不一致的结果 - “一只失去一只腿的熊”。这种视觉比较突显了我们方法在各种任务中的强大和多功能性。

定量评估

用户研究。我们将我们的方法与已发布的ControlNet模型以及重新实现的单任务ControlNet在我们的训练集上进行了性能比较。如图6所示，我们的方法在所有情况下始终优于替代方法。在HED到图像生成任务中，我们的方法明显优于ControlNet。这种优势在深度和法线表面到图像生成任务中更为明显，用户普遍青睐我们的方法，表明其处理复杂几何解释的能力。与重新实现的单任务模型相比，图7显示了我们的方法仍然具有较小的优势，但它仍然通过有效地识别图像区域来引导内容生成来展示其优势。

即使在具有挑战性的Outpainting任务中，我们的模型也优于基线，突显了其鲁棒性和泛化能力。

零样本泛化

我们进一步展示了我们的方法在处理混合条件组合和未见任务泛化的零样本挑战方面的惊人能力。

混合任务组合。这涉及同时从两个不同的条件生成结果。我们的模型的零样本能力通过深度和人体骨架或分割地图和人体骨架等组合进行测试。结果如图5(a)-(b)所示。当背景条件为深度图时，模型有效地描绘了场景的复杂三维结构，同时保持了人体主体的骨骼结构。类似地，当模型面对分割地图和人体骨架的组合时，输出巧妙地保留了主体的结构细节，同时遵循分割边界。这些示例说明了我们模型的适应性和鲁棒性，突出了其在没有任何先前明确训练的情况下处理复杂混合任务的能力。

未见任务泛化。为了评估零样本泛化能力，如灰度图像上色、图像去模糊和图像修复等未见任务，我们在图5©-(e)中进行了案例分析。该模型巧妙地处理了未见任务，产生了引人入胜的结果。这种能力深植于预训练任务和新任务之间的共享属性和隐含相关性，使我们的模型能够无缝地适应。例如，上色任务利用了模型对来自分割任务和深度估计任务的图像结构的理解，而去模糊和修复任务则受益于模型对边缘检测和外扩任务的熟悉。

结论和讨论

我们引入了UniControl，这是一种新颖的统一模型，用于将各种条件纳入扩散模型的生成过程中。UniControl通过采用两个关键组件，即专家混合（MOE）样式适配器和任务感知HyperNet，被设计成可以适应各种任务。实验结果展示了该模型在不同任务和条件下的稳健性能和适应性，表明其在处理复杂的文本到图像生成任务方面具有潜力。

局限性和更广泛的影响。虽然UniControl表现出令人印象深刻的性能，但它仍然继承了基于扩散的图像生成模型的局限性。具体来说，它受到我们的训练数据的限制，这些数据来自Laion-Aesthetics数据集的一个子集。我们观察到这个数据集存在数据偏差。尽管我们已经执行了关键词和基于图像的数据过滤方法，但我们意识到模型可能生成有偏见或低保真度的输出。当需要高质量的人类输出时，我们的模型也存在局限性。如果能够提供更好的开源数据集以阻止生成有偏见、有害、有害性化或其他有害内容，UniControl可能会得到改进。我们希望我们的工作能够激发研究人员开发视觉生成基础模型。