OmniGen: Unified Image Generation 旨在统一常见生成任务比如文生图、图片编辑、图片复原、可控生成，甚至姿态估计等理解任务到同一个模型中。相比于最近很火的理解生成一体化，OmniGen更关心生成，它将以上任务，统一地放在了一个“图文交叉序列输入 + 图片输出”的框架中，训练一个模型实现了上面所说的所有任务，而且还后续微调新增新的功能还很容易。该工作来自于智源，开源了权重和训练代码： VectorSpaceLab/OmniGen。
网络#
OmniGen是一个Diffusion模型，微调自预训练的LLM（Phi3），LLM接受图文交叉序列作为输入：图像经过预训练且fixed的SDXL VAE编码后，2x2的latent patch经过Linear映射后作为图片token embedding，被新引入的<img></img>token包一下后，和文本token embedding拼在一起。
和DiT里in-context的模型一样，Transformer输入端会再继续拼上timestep embedding和当前t时刻的noisy latents。训练时attention mask部分如图所示，相对自然。text部分依然是causal attention mask，image patch内部则是双向注意力。后面的timestep和noise可以看到前面的所有embedding。推理时，text + image token部分可以像LLM一样KVCache住，不需要每个Diffusion step都计算。
插入一些个人想法#
OmniGen的网络结构、输入输出设计非常简洁自然，丢掉了SDXL、FLUX这些模型常见的condition encoder部分，同一个Transformer既做理解，也负责预测噪声。这种设计在网络结构上与VLM很类似，能享受到很多LLM领域的工程优化。比如条件编码部分可以用KVCache，特殊的attention mask也可以用到FlexAttention。
但这种结构，我比较担心“MoE地够不够”。Diffusion训练是一个多任务训练过程，同一个网络需要完成输入输出和预测难度完全不同的预测任务，需要模型“MoE化”，比如DiT最后选用的AdaLN-zero，time embedding直接通过AdaLN-zero修改模型中间层的分布。而OmniGen选取了DiT中证明不太好的in-context time embedding策略，只是在输入加了一个timestep embedding，处理不同难度的Diffusion/Flow预测任务时可能会互相冲突。
这一点有啥改进方法吗？我首先想到的是一个很脏的方法：
timestep embedding之前的token用LLM原本的LN/RMSNorm，之后的token用AdaLN。这样可以训，但是训练时要不断地split然后concat特征，可能会比较慢。
这个思路再进化一下，就是MMDiT，condition部分还用原来LLM的东西，noisy latents token部分则用单独的一套toqkv和FFN，但是和condition一起做attention。相当于把MMDiT结构里固定的text encoder部分省略，从pretrained LLM开始微调。
另外一种思路是，noisy latents不要直接进transformer，而是通过一个简单的带着AdaLN之类东西的transformer层首先处理一下，再变成LLM transformer输入。这里的输入还可以用最近刚出现的REPA，要求和DINOv2特征对齐一下，更语义化，再作为LLM transformer的输入。
训练策略#
OmniGen用rectified flow作为基础生成算法。看了下开源版本的代码，似乎没有time shift相关的东西，但是用了lognorm来采样timestep。作者在论文里专门强调了针对编辑任务需要改变一下损失函数，在空间维度根据编辑前后图片的差异设置了加权，迫使模型更关注目标图片和生成图片不一样的地方。只不过论文里的loss权重公式似乎是写错了，变成了差异越小，权重越大(>1)；没有差异时，权重手工设置为1。
训练根据分辨率分了很多个阶段：
最大分辨率2240还是有些意外，一般模型都没做这么大。104这个batch size，对应的单卡bs为1，应该还有优化空间，一共不到4B的模型，应该能训练更快的。
数据#
接下来来看看最关键的训练数据。OmniGen把自己的训练数据集叫做"X2I"，指anything to image。整个数据集包含100M样本。每个样本都被统一为“图文交叉序列->图片”的形式，输入端训练时最多有3个图片。
| 分类 | 子任务 | 数据集 |
|---|---|---|
| Text to Image Generation | 文生图使用了下面这些数据和16M的内部数据（猜测是爬的MJ、pexels之类的数据）。用InternVL2给内部数据和LAION-Aesthetic做了recaption。 1. Recap-DataComp的子集，约56M images. 2. SAM-LLaVA 3. ShareGPT4V 4. LAION-Aesthetic子集，约4M 5. ALLaVA-4V 6. DOCCI 7. DenseFusion 8. JourneyDB |
|
| Common Mixed-modal Prompts | image editing | SEED-Data-Edit, MagicBrush, and InstructPix2Pix |
| human motion | Something-Something | |
| virtual try-on | HR-VITON and FashionTryon | |
| spatial control | MultiGen：Canny, HED, Depth, Skeleton, Bounding Box, and segmentation. | |
| style transfer | stylebooth | |
| Subject-Driven Image Generation | 图片中object crop recontext | GRIT-Entity dataset. 分割出GRIT中的物体，然后inpainting生成物体的背景。相当于给原图中物体换背景后当做condition。6 million pairs. |
| 名人数据集 | 选了2000个名人和这些人相关的人，共10000个人，搜索相关图片，用InternVL删除搜索结果中非相关的图片。构建了533K pair的数据集，这里的pair可能是指<personA, personB, 同时包含personA和personB>这样的三元组数据集。 | |
| Computer Vision Tasks | low-level vision tasks | low-light image enhancement, deraining, deblurring, inpainting, outpainting, and colorization 输入图片 |
| human pose, depth mapping, canny, and segmentation | LAION作为source 图，类似于UniControl，用专家模型制作这些图片的对应任务标注。 | |
| referring image segmentation | RefCOCO, ADE20k and ReasonSeg。把需要指出的区域用蓝色区域高亮出来，当做目标生成图 | |
| Few-shot to Image | 上面的任务，随机选另外一个样本pair当做few-shot案例，让模型照着预测当前这个样本输入图对应的生成图。模拟学习in-context learning capabilities |
subject-driven image generation部分的数据构造的pipeline如下：
可惜的是，论文没有透露太多这些不同任务的数据的配比。训练时不同难度的任务是否需要微调训练顺序等等。
实验#
论文首先展示了一些OmniGen能做的任务：
文生图语义相关的指标也很高：
作者常规地测试了很多训练中的任务的效果，看评测结果都还不错。
我们主要关注OmniGen这种anything2image的混合多任务训练，有没有让模型具备一些训练集中没有的新能力：
作者首先展示了OmniGen可以一步实现参考图片的depth/pose生成target图这样的任务。既原本需要首先用理解模型推理出中间pose/depth图，然后再作为condition，生成target图。而OmniGen可以根据prompt，一步完成参考图片中的pose或者depth或者Canny edge，生成新图片的任务。
这项能力，应该来自于模型训练时的spatial control任务和computer vision任务能力的组合。
作者还展示了模型In-context Learning for Unseen Tasks and Domains的能力。比如下面最后的分割任务，直接让模型做是做不了的，但是给它一个示例sample，就可以做了。说明模型具有一定的in-content learning的能力。
OmniGen还展示了模型的推理能力，比如可以找到图片中可以洗手的地方。这个点有点怪，这个能力应该是ReasonSeg这个训练集里的一部分吧。不太能说是模型新涌现出来的能力。
作者还用张吕敏的PAINTS-UNDO构造了渐进式画图的数据集，微调了OmniGen，试图展示OmniGen实现Chain of Thought的潜力。他们试图模拟人类画图时的过程，先画一个初稿，然后将初稿再作为模型输入生成下一个环节的图片，重复多次后最终画出目标图片。
作者提到，根据LLM领域的经验，对过程监督好于只对结果监督。他们期望这样一步一步作画的能力，能启发后续图片生成领域对过程监督的工作。这个思路确实挺有趣的，Diffusion本身实际上已经可以看作是对画图过程进行监督的工作了（监督低频、然后监督高频…)。如果我们能更进一步，从单张图片的高低频，变成类似于人的对底稿进行监督，对单个物体进行监督，对细节进行监督…，是否能获得一些更高级的能力？
我想，也许把人类PS数据当做训练集也是个思路。比如选有畸变的图片，人工修复后，分别当做OmniGen的输入输出，微调模型。这样，每次推理后，再过模型refine一下，能修复一下常见的肢体扭曲啥的。这也算分解生成任务了。
总结#
我一直对图文交叉序列->图片这种范式情有独钟。很多常见的图片生成领域业务问题实际上都是这个任务下的一个特例。有了这样的基础模型，加不同数据，就能微调出各个业务需要的模型。而不是每个任务都专门设计半天IP-Adapter。
接下来等我首先测试下OmniGen的效果，看看到底到了什么层面，上限高不高，有没有机会达到类似于FLUX这样的生成能力。
后续，也许可以观察下以下几个方向的工作：
- OmniGen训练时完全放弃了LLM本身的语言能力，怎么设计，能让模型保留文字生成的能力？也不需要完全留住LLM的能力，可以在生成图片前先大概描述一下要生成的图片的内容或者需要编辑的区域，这样是否也能实现“对画图过程的监督”。
- 可能保留FLUX或者SD3.5的生成能力，但把它们的condition部分，替换成类似于OmniGen的图文交叉序列吗？