读论文《OmniGen: Unified Image Generation》

OmniGen演示了在单一框架内执行各种图像生成任务的能力。此外，它还具有推理能力和语境学习能力。

论文地址：2409.11340v1 (arxiv.org)

项目地址：GitHub - VectorSpaceLab/OmniGen

项目目前还不完整，论文展现的通用性十分强大，就是不知道稳定性怎么样。

摘要

大型语言模型（LLM）的出现统一了语言生成任务，并彻底改变了人机交互。然而，在图像生成领域，能够在单个框架内处理各种任务的统一模型在很大程度上仍未得到探索。在这项工作中，我们介绍了 OmniGen，这是一种用于统一图像生成的新扩散模型。与流行的扩散模型（例如 Stable Diffusion）不同，OmniGen 不再需要 ControlNet 或 IP 适配器等额外模块来处理各种控制条件。OmniGen 具有以下特点：1）统一：OmniGen 不仅展示了文本到图像的生成功能，而且本身也支持各种下游任务，例如图像编辑、主题驱动生成和视觉条件生成。此外，OmniGen 可以通过将经典的计算机视觉任务转换为图像生成任务（例如边缘检测和人体姿势识别）来处理这些任务。2）简单性：OmniGen 的架构高度简化，无需额外的文本编码器。此外，与现有的扩散模型相比，它更加用户友好，能够通过指令完成复杂的任务，而无需额外的预处理步骤（例如，人体姿态估计）并花费大量成本，从而显著简化了图像生成的工作流程。3）知识转移：受益于统一格式的学习，OmniGen 可以有效地在不同任务之间传递知识，管理看不见的任务和领域，并展示出新颖的能力。我们还探讨了该模型的推理能力和思维链机制的潜在应用。

统一：不仅可以文本生成图像，还支持多种下游任务，如图像编辑、主题驱动生成和视觉条件生成。（这个准确率真的可以保证吗）
简单：架构简化，不需要额外的文本编码器。
知识迁移：通过统一格式学习，有效跨任务迁移知识，处理未见任务和领域。

OmniGen

原则：1）通用性：接受任何形式的图像和文本输入以完成各种任务;2）简洁，避免过于复杂的结构设计和众多的附加组件。

网络架构：如图 2 所示，OmniGen 框架采用由变分自动编码器（VAE）和预训练的大型转换器模型组成的架构。具体来说，VAE 从图像中提取连续的视觉特征，而 transformer 模型根据输入条件生成图像。在本文中，我们使用 SDXL的 VAE 并在训练期间将其冻结。我们使用 Phi-3 [1] 来初始化 transformer 模型，继承了它出色的文本处理能力。与需要额外编码器来预处理条件信息（例如剪辑文本编码器和图像编码器）的最先进的扩散模型不同，OmniGen 本身就对条件信息进行编码，从而大大简化了管道。此外，OmniGen 在单个模型中联合建模文本和图像，而不是像现有工作 [67; 68; 70; 63; 9] 那样使用单独的编码器独立建模不同的输入条件，这缺乏不同模态条件之间的交互。

OmniGen 的框架。文本被标记化为标记，而输入图像通过 VAE 转换为嵌入。OmniGen 可以接受自由格式的多模态提示，并通过整流方法生成图像。

输入格式：模型的输入可以是自由形式的多模态交错文本和图像。我们使用 Phi-3 的分词器来处理文本，无需任何修改。对于图像，我们首先使用具有简单线性层的 VAE 来提取潜在表示。通过将每个补丁线性嵌入潜在空间中，将其转换为视觉标记序列。按照 [50]，我们将基于标准频率的位置嵌入应用于输入视觉标记，并使用与 SD3 [13] 相同的方法来处理具有不同纵横比的图像。此外，我们在将每个图像序列插入 t 之前，用两个特殊标记封装每个图像序列：“<img>”和“</img>”

注意力机制：修改了 LLM 中常见的因果注意力机制，将其与双向注意力集成，如图 2 所示。具体来说，我们对序列中的每个元素应用因果注意力，但在每个图像序列中应用双向注意力。这允许每个 patch 注意同一图像中的其他 patch，同时确保每个图像只能关注以前出现的其他图像或文本序列。

推理：在推理过程中，我们随机采样一个高斯噪声，然后应用流匹配方法来预测目标速度，迭代多个步骤以获得最终的潜在表示。最后，我们使用 VAE 将潜在表示解码为预测图像。默认推理步骤设置为 50。得益于注意力机制，OmniGen 可以像使用 kv-cache 一样加速推理，就像 LLMsby：存储前一个和当前的 key 和 va

训练策略

训练目标：使用修正的流（rectified flow）来优化模型参数。模型被训练以直接回归给定噪声数据、时间步和条件信息的目标速度。
训练管道：逐渐增加图像分辨率，使用 AdamW 作为优化器。

X2IDataset

在论文中，X2I数据集是为了训练和评估 OmniGen 模型而构建的大规模统一图像生成数据集。X2I 的含义是 "anything to image"，即任何内容到图像，这个数据集涵盖了多种图像生成任务，并且将它们统一成一个格式以供模型学习。

1. 数据集构建目的

X2I 数据集旨在提供多样化的任务和数据类型，以支持 OmniGen 模型的训练，使其能够处理各种图像生成任务。这包括但不限于文本到图像的生成、图像编辑、主题驱动的图像生成以及视觉条件生成。

2. 数据集组成

X2I 数据集由以下几个部分组成：

2.1 文本到图像（Text to Image）

数据来源：包括 Recap-DataComp、SAM-LLaVA、ShareGPT4V、LAION-Aesthetic、ALLaVA-4V、DOCCI、DenseFusion 和 JourneyDB 等。
数据使用：早期训练阶段使用这些数据集来学习广泛的图像-文本匹配关系和多样化的知识。

2.2 多模态到图像（Multi-modal to Image）

这一部分的数据输入是任意交错的文本和图像。

混合模态提示：例如图像编辑、人类动作、虚拟试穿和风格迁移等任务。
视觉条件：使用 MultiGen 数据集学习更细粒度的空间控制功能，包括边缘检测、深度图、骨架、包围盒和分割等六种代表性视觉条件。

2.3 主题驱动图像生成（Subject-Driven Image Generation）

基础数据集：GRIT-Entity 数据集，通过使用 GRIT 数据集的标注，结合 Grounding DINO 模型和 SAM 进行对象检测和分割，然后使用 MS-Diffusion 模型进行图像重绘。
高级数据集：Web Images 数据集，使用自然图像中的知名个体，通过交叉验证策略筛选和标注图像。