如您有工作需要分享，欢迎联系：aigc_to_future
作者：Chong Mou等
解读：AI生成未来
文章链接：https://arxiv.org/pdf/2504.16915
项目链接：https://mc-e.github.io/project/DreamO
亮点直击
提出了 DreamO，一个统一的图像定制框架。它通过在预训练的 DiT 模型上训练少量额外参数，实现了各种复杂和多条件的定制任务。 基于扩散模型内部的表示对应关系，本文设计了一种特征路由约束，以增强一致性保真度，并在多条件场景中实现有效的解耦。 引入了一种渐进式训练策略，以促进多任务和复杂任务设置的收敛。本文还设计了一种占位符策略，以建立文本描述与条件图像之间的对应关系。 大量实验表明，本文的方法不仅在广泛的图像定制场景中实现了高质量的结果，而且在适应多条件场景时表现出强大的灵活性。
总结速览
解决的问题
任务局限性：现有图像定制方法多为特定任务设计（如身份保持、风格迁移、虚拟试穿等），缺乏通用性，难以整合多类型控制条件。 多条件交互困难：现有统一框架（如UniControl、OminiControl）对多条件的组合控制能力不足，存在冗余或交互僵化问题。 训练成本高：部分方案（如OmniGen、UniReal）依赖大语言模型或视频预训练，需数百GPU资源，成本高昂。
提出的方案
统一框架设计：基于扩散Transformer（DiT）构建通用架构，通过统一序列处理多类型输入（如身份、物体、风格等）。 特征路由约束：在训练中引入特征路由机制，精准查询参考图像信息，提升内容保真度并解耦不同控制条件。 占位符策略：将条件与生成结果中的特定位置关联，实现条件放置的灵活控制。 渐进式训练策略： 初期：简单任务训练，建立基础一致性。 中期：全规模训练，综合提升定制能力。 后期：质量对齐阶段，修正低质量数据偏差。
应用的技术
扩散Transformer（DiT）：作为主干模型，统一处理多模态输入条件。 统一序列编码：将不同条件（如草图、深度图、参考图像）转换为序列化输入。 特征路由约束：通过注意力机制优化条件特征的检索与融合。 大规模多任务数据集：构建涵盖身份、物体、风格等多样化任务的数据集。
达到的效果
多功能支持：单模型支持身份保持、物体外观迁移、虚拟试穿、风格化等多种任务。 高质量生成：通过渐进训练和质量对齐，生成结果保真度高且无明显偏差。 灵活控制：占位符策略实现多条件的精确位置控制，支持复杂场景生成（如多物体插入）。 高效训练：仅需轻微训练成本（对比OmniGen等方案），即可实现多条件交互与统一部署。
方法
概述
本文的方法概览如下图2所示。本文采用Flux-1.0-dev作为基础模型，构建一个支持风格、身份、主体外观和虚拟试穿等任务的统一图像定制框架。给定n个条件图像，首先复用Flux的VAE将条件图像编码到与噪声隐空间变量相同的隐空间。随后，通过分块化（patchification）将2D隐空间表示转换为1D token序列，所有tokens沿序列维度拼接后输入Flux。为使模型能融合条件图像输入，在Flux的输入端引入专用映射层来编码条件图像tokens。这些条件tokens的位置嵌入（PE）通过旋转位置嵌入（RoPE）与噪声隐空间变量的位置对齐。受OminiControl非重叠位置嵌入的启发，沿对角线以类似方式扩展这些嵌入。为区分隐空间中的条件隐空间变量和噪声图像隐空间变量，在输入阶段向条件隐空间变量添加可学习的条件嵌入（CE）。此外，为支持多条件场景，我们引入可学习的索引嵌入（IE），并根据条件索引将其添加到每个条件隐空间变量中。遵循OminiControl的设计，我们将低秩自适应（LoRA）模块集成到Flux中作为可训练参数，使模型能更好地适应条件任务。
路由约束
受UniPortrait和AnyStory的启发，本文在DiT架构中设计了适用于通用图像定制任务的路由约束。如前面图2所示，在条件引导框架中，条件图像与生成结果之间存在交叉注意力：
其中 表示第 个条件图像的tokens序列， 是噪声图像隐空间变量的tokens序列。交叉注意力图 表示第 个条件图像与生成结果之间的密集相似度关系。为获取条件图像在生成结果不同位置上的全局响应，沿 维度对密集相似度矩阵进行平均，得到响应图 ，该图反映了条件图像对生成结果的全局相似性响应。
为可视化 DiT 框架中的图像间交叉注意力图，进行了简化实验。具体而言，使用 Subject200k 数据集在图 2 的可控框架上训练主体驱动生成任务。下图 3 第一行展示了生成结果及内部交叉注意力图，本文选择 Flux 双块结构中的三个层级（第 2、8、14 层）进行注意力可视化。可以观察到：在 DiT 框架下，图像间注意力呈现空间模式，虽粒度较粗但能聚焦目标物体。
为约束图像间注意力集中于特定主体，采用 MSE 损失优化 DiT 中条件图像与生成结果间的注意力关系：
其中 和 分别表示条件索引和层索引， 和 分别为条件数量和层数。 表示目标图像中的主体掩码。
除图像间路由约束外，还设计了占位符-图像路由约束以建立文本描述与条件输入的对应关系。具体而言，对于第 个条件，我们在对应实例名称后添加占位符 [ref#i]（例如："A woman from [ref#1] and a woman from [ref#2] is walking in the park"）。在多条件任务训练期间，我们计算条件图像tokens与占位符tokens之间的相似度。该路由约束确保：
条件图像tokens与自身占位符 [ref#i]的相似度为 1与其他占位符的相似度为 0
表示为：
其中 表示占位符 [ref#i] 的特征， 是一个二元矩阵（当占位符与条件图像匹配时值为1，否则为0）。
本文方法的最终损失函数定义为：
其中，、 和 是损失权重。需要注意的是，为了让模型能够适应常规文本输入，我们在训练时以50%的概率引入不含占位符的普通文本，并相应地舍弃 。
如上图3第二行所示，经过路由约束训练后，条件图像的注意力明显聚焦于目标主体，生成结果在细节上与参考图像的一致性也有所提升。除了增强一致性外，该策略还有助于多参考情况下的解耦。更多细节见消融实验（下图11）。下图4展示了占位符的效果，它可以控制生成结果中条件的位置。
训练数据构建
本工作旨在实现广义的图像定制。本文收集了一个覆盖广泛任务的训练数据集。
身份配对数据
由于高质量的身份配对数据难以从互联网获取，我们采用开源ID定制方法PuLID进行数据集构建，因其展现出高面部相似性和强提示对齐。具体而言，我们为PuLID-SDXL提供参考人脸图像和描述目标风格的文本提示，使其控制生成肖像的风格，从而得到（参考人脸，提示，风格化人脸）形式的训练对。对于写实场景，我们使用PuLID-FLUX生成同一身份的两张图像作为相互参考。
主体驱动数据
对于单主体驱动的图像定制，本文使用开源数据集Subject200K作为训练数据。为补充角色相关条件的缺失，我们通过检索和人工筛选收集了一组成对的角色相关数据。对于多主体驱动的图像定制，我们在Subject200K数据集上通过图像拼接构建双列图像。此时，文本提示变为："生成一张双列图像。左侧为{prompt1}，右侧为{prompt2}"，其中prompt1和prompt2分别对应两幅拼接图像的描述。为增强处理复杂场景（尤其是涉及同一类型的多个参考因素）的能力，我们在通过图像拼接构建多主体图像时，会按一定比例从同一类别中选择主体。本文还使用开源数据集X2I-subject进行多主体驱动训练。
为增强以人类为参考的主体驱动生成，仿照MoiveGen的流程构建了另一个数据集。从长视频数据集出发，通过内容感知的场景检测获取多个短视频片段。为获取片段内实例匹配，我们使用Mask2Former提取关键帧中所有人的掩码，并在视频片段中进行目标跟踪。为获取跨片段实例匹配，本文采用SigLip嵌入所有人体对象并进行聚类。之后，使用GPT-4o结合图像和彩色实例掩码输入生成实例 grounded caption。
试穿数据
在该任务中，通过两种来源构建成对的试穿数据集。部分数据直接从网络收集模特与服装的配对图像；另一部分则先爬取高质量模特图像作为真值，再通过图像分割提取服装并构建对应配对数据。所有收集的图像均经过人工过滤以剔除低质量样本。服装区域通过分割算法提取。
风格驱动数据
本文旨在解决两类风格迁移任务：(1) 风格参考图像+内容文本描述控制；(2) 风格参考图像+内容参考图像。对于第一类任务，训练数据需包含同一风格的参考图像与目标图像。如图13所示，我们采用内部风格迁移模型（即基于SDXL的模型，使用类似InstantStyle的解耦策略），在两种不同提示下生成相同风格但内容各异的图像。对于第二类任务，训练需要风格参考图像、内容参考图像和目标图像。此时目标图像与风格参考共享风格，同时保持与内容参考图像相同的内容结构。基于第一类训练数据，我们通过Canny-guided Flux生成与风格图像对应的自然图像作为内容参考。
路由掩码提取
为了获取路由约束（即）的标签，我们从各种类型的训练数据中提取目标主体掩码。具体来说，对于X2I-Subject数据集，使用InternVL模型提取主体描述，然后利用LISA结合描述和目标图像生成主体掩码。对于Subject200K数据集，使用InternVL获取主体名称，然后通过LISA预测相应的掩码。对于人脸数据，我们直接使用LISA从目标图像中提取人脸掩码。
尽管路由约束仅应用于有限的训练数据（即Subject200K、人脸数据和X2I-Subject），但本文的方法在未见场景中表现出强大的泛化能力。如下图5所示，第一行显示在试穿任务中，条件图像准确地定位了目标区域。第二行显示在更复杂的条件下，模型有效地区分了与不同条件图像相关的区域。具体来说，身份参考图像聚焦于面部和身体，而带有太阳镜的图像则直接关注太阳镜区域，避开了面部。
渐进式训练流程
在实验中，本文发现直接在全部数据上训练会导致收敛困难。这主要是由于优化参数的容量有限，使得模型在复杂数据分布下难以捕获特定任务的能力。此外，由于训练数据图像质量的影响，训练后的生成质量与Flux的生成先验存在偏差。
为了解决这些问题，本文设计了一种渐进式训练策略，使模型能够在不同任务间平滑收敛，同时纠正训练数据对Flux高质量生成先验的影响。训练流程如下图6所示。首先在主体驱动的训练数据上优化模型，以初始化模型的一致性保持能力。需要注意的是，训练中使用的Subject200K数据集是由基模型（即Flux）生成的，因此与模型生成空间具有相似的分布，这有助于快速收敛。由于X2I-subject数据集是通过MS-Diffusion合成的，许多训练样本包含不希望的伪影和失真。因此，在这个预热阶段，还使用了双列Subject200K图像作为训练数据的一部分，以促进多主体生成控制的快速收敛。图6的右侧部分显示，在第一个训练阶段后，模型获得了初步的主体驱动生成能力，并表现出强大的文本跟随性能。在第二个训练阶段，我们整合所有训练数据并进行全数据微调，使模型能够进一步收敛到本工作中定义的所有子任务上。
在全数据训练的第二阶段后，生成质量受到训练数据的严重影响，尤其是低质量的训练样本。为了使生成质量与Flux的生成先验重新对齐，本文设计了一个图像质量精炼训练阶段。具体来说，我们利用Flux生成了约40K个训练样本。在训练过程中，使用原始图像作为参考来指导模型重建自身。为了防止复制粘贴效应，丢弃了95%的参考图像token。经过短时间优化后，生成质量显著提高，实现了与Flux生成先验的对齐。
实验
实现细节
本文采用Flux-1.0-dev作为基础模型，设置额外LoRA的秩为128，共增加707M参数。训练过程使用Adam优化器，学习率设为4e-5，在8块NVIDIA A100 80G GPU上进行训练。批大小设置为8。第一阶段训练20K次迭代，第二阶段90K次迭代，最后阶段3K次迭代。
模型能力
DreamO是一个统一框架，能够处理多种图像定制任务。下图7展示了身份驱动的图像生成，本文的方法既支持特定个体的定制，也支持组合身份控制。模型可以生成保留单个人物身份或融合多个人物特征的图像，同时通过文本输入精确控制其他属性和场景细节。
下图8展示了物体驱动的图像定制结果，可见DreamO既支持单主体条件控制，也支持多主体组合，能够将多样化视觉元素整合到统一输出中并保持语义一致性。
下图9展示了试穿能力，包括特定身份的视觉试穿和任意主体的创意应用。
下图10展示了风格驱动定制，风格完全由参考图像引导，其他属性通过文本或图像输入控制，实现细粒度的多模态操控。这些结果凸显了本文方法整合异质控制信号（如身份、物体和试穿）的能力，实现更具表现力和复杂度的图像定制。
消融研究
路由约束：在DiT训练中引入路由约束以提升生成保真度并促进多条件解耦。下图11显示，在单条件生成中去除该约束会导致参考保真度下降（如服装颜色与参考不一致）；在多条件设置中会导致条件耦合（如两个玩具特征交叉）。这些结果证实路由约束能提升不同条件的保真度和解耦性。
渐进式训练：为使模型在复杂数据分布下更好收敛并修正训练数据分布对生成质量的影响，设计了渐进式训练策略。图12表明，直接在全部数据集上训练会导致次优收敛（尤其在多主体一致性等复杂任务中）；先在较小且易学的数据集（如Subject200K）上预热再联合训练可改善收敛，但生成质量易受训练数据分布影响而偏离Flux的生成先验；通过引入图像质量调优阶段，模型能产生更高质量的生成结果。
结论
本文提出了DreamO，一个基于预训练DiT框架的统一图像定制系统，能够处理多种控制条件（如身份、风格、主体和试穿）。首先构建了包含广泛图像定制任务的大规模训练数据集。通过将所有条件类型整合到DiT输入序列，并引入特征路由约束，DreamO在确保高保真一致性的同时，有效解耦了异构控制信号。
针对复杂数据分布下保持基础模型（Flux）生成先验的同时学习强任务特定能力的挑战，设计了渐进式训练策略。该方法使模型能逐步获得多样化控制能力，同时保持基础模型的图像质量。大量实验表明，DreamO能够高质量地完成各类图像定制任务。此外，其基于LoRA的轻量级设计实现了低计算成本的高效部署。
参考文献
[1] DreamO: A Unified Framework for Image Customization
致谢
如果您觉得这篇文章对你有帮助或启发，请不吝点赞、在看、转发，让更多人受益。同时，欢迎给个星标⭐，以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进，共同探索未知，见证一个充满希望和伟大的未来！
技术交流
加入「AI生成未来社区」群聊，一起交流讨论，涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向，备注不同方向邀请入群！可添加小助手备注方向加群！
没有评论:
发表评论