Source-linked AI summary
OmniGen: Unified Image Generation
Shitao Xiao, Yueze Wang, Junjie Zhou, Huaying Yuan, Xingrun Xing, Ruiran Yan, Chaofan Li, Shuting Wang, Tiejun Huang, Zheng Liu
TL;DR
Image generation lacks a unified model that handles diverse tasks within one framework. OmniGen introduces an instruction-driven diffusion model trained on a unified multi-task dataset, reporting broad capabilities across tasks while remaining simple and end-to-end; the authors also identify limitations in text rendering, undesired details, and unseen image types.
Problem
Image generation lacks a counterpart to LLMs that can handle diverse tasks within a single framework.
Method
OmniGen combines a concise VAE-transformer architecture with unified multi-task training on X2I and rectified-flow optimization.
Results
OmniGen demonstrates image-generation capabilities across varied tasks, including downstream vision tasks and complex end-to-end instruction following.
Takeaways & Limitations
Unified training allows OmniGen to transfer learned knowledge to unseen tasks and domains and exhibit new capabilities.
Takeaways & Limitations
Text rendering remains limited, undesired details may appear, and previously unseen image types such as surface normal maps can be difficult to process.
Abstract
from arXiv · showhide
The emergence of Large Language Models (LLMs) has unified language generation tasks and revolutionized human-machine interaction. However, in the realm of image generation, a unified model capable of handling various tasks within a single framework remains largely unexplored. In this work, we introduce OmniGen, a new diffusion model for unified image generation. OmniGen is characterized by the following features: 1) Unification: OmniGen not only demonstrates text-to-image generation capabilities but also inherently supports various downstream tasks, such as image editing, subject-driven generation, and visual-conditional generation. 2) Simplicity: The architecture of OmniGen is highly simplified, eliminating the need for additional plugins. Moreover, compared to existing diffusion models, it is more user-friendly and can complete complex tasks end-to-end through instructions without the need for extra intermediate steps, greatly simplifying the image generation workflow. 3) Knowledge Transfer: Benefit from learning in a unified format, OmniGen effectively transfers knowledge across different tasks, manages unseen tasks and domains, and exhibits novel capabilities. We also explore the model's reasoning capabilities and potential applications of the chain-of-thought mechanism. This work represents the first attempt at a general-purpose image generation model, and we will release our resources at https://github.com/VectorSpaceLab/OmniGen to foster future advancements.
1. Introduction
OmniGen addresses the lack of a universal image-generation counterpart to LLMs by unifying multiple tasks in one instruction-driven framework. It combines a concise architecture and unified training format to support flexible, end-to-end image generation.
- Motivation: Current image-generation models specialize in individual tasks and often require task-specific modules and fine-tuning.This limits a single model’s ability to handle diverse tasks through one workflow.
- Motivation: OmniGen asks whether one model can complete varied image-generation tasks end-to-end through user instructions, without additional plugins or operations.The framework is intended to make image generation simple and flexible for arbitrary instructions.
- Framework: OmniGen uses a concise VAE-and-transformer architecture that accepts arbitrarily interleaved text and image conditions.Its multi-modal input design extends beyond pure-text or pure-image conditioning.
- Unified training: X2I unifies diverse image-generation tasks in a common format for training OmniGen as a general-purpose model.The dataset is presented as a comprehensive resource spanning multiple image-generation tasks.
- Unified training: Unified training enables OmniGen to apply learned knowledge to unseen tasks and domains while exhibiting new capabilities.The paper also explores reasoning and chain-of-thought mechanisms.
2. OmniGen
OmniGen combines a VAE with a transformer, represents free-form interleaved text and images as a unified sequence, and trains with rectified flow. Its training strategy also emphasizes changed regions in image editing.
- Architecture: OmniGen comprises a VAE and a pre-trained transformer, with the VAE extracting visual features and the transformer generating images from instructions.Only the VAE is frozen during training.
- Input representation: Free-form inputs interleave text tokens and image tokens, with image latents flattened into visual-token sequences and no task-specific special tokens.Images are encoded with a VAE and linear layers before sequence insertion.
- Attention mechanism: The attention mechanism applies causal attention across the sequence and bidirectional attention within each image sequence.This lets patches attend to other patches in the same image while preserving cross-image and text sequencing constraints.
- Training objective: Rectified flow linearly interpolates between Gaussian noise and data, training the model to regress target velocity from noised data, timestep, and conditions.The objective minimizes mean squared error between predicted and target velocity.
- Image editing: Image-editing loss weights are amplified in altered regions to prevent the model from minimizing loss by simply copying the input image.Regions with changes receive significantly higher weights than unchanged regions.
- Training pipeline: Training progressively increases image resolution because lower resolution is data-efficient while higher resolution can improve aesthetic quality.The pipeline uses AdamW and runs experiments on 104 A800 GPUs.
3. X2I Dataset
X2I is a large-scale unified dataset that standardizes diverse image-generation tasks as interleaved image-text inputs. It covers text-, image-, subject-, few-shot-, and computer-vision-to-image tasks.
- X2I contains approximately 0.1 billion images converted into a unified format for robust multi-task image generation.
- The dataset combines plain-text image descriptions with arbitrarily interleaved text and image inputs from editing, motion, virtual try-on, style transfer, and visual-condition tasks.
- Subject-driven data pairs reference and target images, using object annotations, grounding, segmentation, and web-collected images of the same person.
- Computer Vision Tasks: Computer-vision tasks include low-level restoration and transformation tasks, plus high-level annotations represented as images for pose, depth, edges, and segmentation.
- Few-shot to Image: Few-shot-to-image inputs combine original task inputs with randomly selected examples, but training uses only one example because of resource limitations.
4. Experimental Results
OmniGen handles text-to-image, editing, subject-driven, visual-conditional, and computer-vision tasks through multimodal instructions. Across reported evaluations, it achieves comparable or competitive performance while simplifying several workflows.
- Figure 4 shows OmniGen handling text-to-image, image editing, and visual-conditional generation from multimodal instructions.
- Subject-driven generation: OmniGen selects required objects from multi-object reference images using textual instructions, without manual cropping or face-recognition preprocessing.
- Text-to-image: OmniGen achieves similar GenEval performance to diffusion models such as SD3 with 3.8 billion parameters versus SD3’s 12.7 billion.
- Editing and subject-driven generation: On EMU-Edit and DreamBench, OmniGen shows performance comparable to current state-of-the-art or best proprietary models, using CLIP-T and CLIP-I measures.
- Visual-conditional generation: OmniGen achieves competitive results across visual conditions, evaluated by similarity between input conditions and conditions extracted from generated images.
- Computer Vision Tasks: OmniGen handles low-level tasks such as deraining, deblurring, and inpainting, as well as high-level tasks such as human-pose recognition.
- Scope boundary: The authors state that OmniGen’s computer-vision capabilities are not intended to surpass specialized state-of-the-art models developed over long periods.
- Workflow comparison: Unlike ControlNet’s detector-and-module workflow, OmniGen completes visual-conditional generation in a single step within one model.
5. Further Analysis
OmniGen exhibits knowledge transfer, task composition, in-context learning, and reasoning capabilities, while exploratory chain-of-thought image generation remains limited in quality.
- Emerging Capabilities: Training on the unified X2I format enables knowledge transfer across tasks and scenarios, including unseen tasks and domains.
- Task Composition: OmniGen can process multiple simultaneous instructions spanning different tasks or repeated instances of the same task.
- In-context Learning for Unseen Domains: With an example, OmniGen can complete a task in an unseen domain despite initially failing to recognize unfamiliar concepts such as pencil sharpeners.
- Reasoning Ability: For underspecified instructions, OmniGen can infer the relevant object from image contents and identify its location, such as a sink for washing hands.
- Chain of Thought: Step-by-step generation lets users iteratively control outputs, but the fine-tuned model does not surpass the original model and may introduce erroneous modifications.
- Chain of Thought: Supervising the image-drawing process is proposed as a promising direction for handling more complex and diverse scenes.
6. Ablation Study
The ablations show that modified attention and weighted loss are important for generation quality and image editing, while VAE and CLIP2 input representations perform similarly.
- Attention Mask: Removing the modified attention produces noisy and distorted images, indicating that unidirectional LLM attention is unsuitable for the rectified flow.
- Weighted Loss: Without weighted loss, image editing encourages copying the input image directly because edited regions are usually small.
- Weighted Loss: Weighted loss prevents the shortcut of copying the input image instead of modifying the intended region.
- Input Image Representation: VAE and CLIP2 provide similar input-image embedding results, while VAE has a slight image-similarity advantage and avoids additional modules.
7. Related Work
Prior image-generation systems are largely task-specific or lack image generation, whereas OmniGen is presented as natively supporting multiple image-generation tasks without preprocessing or assistance from other models.
- Generative Foundation Models: Language and multimodal language models learn broad tasks, but multimodal models generally lack the capability to generate images.
- Diffusion Model: Diffusion models such as Stable Diffusion, DALL-E, and Imagen are predominantly designed for text-to-image generation.
- Diffusion Model: Many instruction-based image-editing methods extend Stable Diffusion through task-specific architectural modifications.
- Diffusion Model: OmniGen natively supports various image-generation tasks without preprocessing steps or assistance from other models.
8. Limitations and Conclusion
OmniGen is introduced as a unified image-generation model supported by the X2I dataset, but it retains limitations in text rendering, unwanted details, and previously unseen image types.
- Conclusion: OmniGen is presented as the first unified image-generation model, with X2I constructed as a large-scale unified image-generation dataset.
- Limitations: Text rendering is limited, long text cannot be generated accurately, and outputs may contain undesired details such as abnormal hands.
- Limitations: Previously unseen image types, including surface normal maps, can hardly be processed as expected.
- Conclusion: The authors envision direct generation from multimodal instructions without complex workflows and position OmniGen as an early step toward universal image generation.
9. More Qualitative Results
OmniGen delivers competitive or superior qualitative results across text-to-image, editing, subject-driven, identity-preserving, and multi-image tasks while supporting flexible multimodal control. The section also documents limitations in text rendering, input-image length, fine details, unseen image types, and image quality.
- Overall results: OmniGen achieves competitive or superior results to task-specific open-source models despite being a general-purpose model.It follows multimodal instructions for complex image-generation tasks.
- Image editing: OmniGen accurately modifies specified targets without altering other content and can execute multiple or inferred-target editing instructions in one step.Unified text-and-image processing supports extensive interaction between modalities.
- Subject-driven generation: OmniGen achieves better subject and text fidelity than Kosmos-G on DreamBench without fine-tuning for the specific object.It preserves reference objects while following new textual instructions.
- Identity-Preserving Generation: OmniGen outperforms InstantID and is comparable to PULID-FLUX 3, while preserving or changing clothing from the input image without face detectors or encoders.Its facial similarity remains below PULID-FLUX, a specialized model.
- Flexible control: OmniGen can extract specified people or clothing from images and process multiple images simultaneously.The section contrasts these capabilities with existing models’ inability to handle multiple objects or inputs.
- Limitations: OmniGen remains sensitive to prompt detail, has limited long-text rendering, accepts at most three training-time input images, and cannot process unseen image types.Generated images may also contain erroneous small details, hands, or facial features.
- Practical implications: Fine-tuning OmniGen for downstream tasks is simpler because multimodal inputs reduce the need for task-specific networks or plugins.The authors give ControlNet and style encoders as examples of components that are unnecessary.
- Limitations: OmniGen’s image quality is lower than the latest FLUX model because of limited high-quality data and insufficient model size.The authors identify expanding data and model size as important directions.
10. Subject-driven dataset
The section presents qualitative examples and describes dataset-construction procedures for subject-driven generation. These materials illustrate flexible image conditioning, while the construction process addresses copy-paste patterns and verifies web-image identities.
- Subject-driven dataset: Directly extracted GRIT data can encourage simple copy-paste patterns, motivating a high-quality web-image dataset for subject-driven generation.
- Qualitative capabilities: OmniGen generates text-to-image outputs with arbitrary aspect ratios.
- Qualitative capabilities: Qualitative examples show image editing, subject fidelity, identity preservation, object extraction, and simultaneous processing of multiple images.
- Dataset construction: The GRIT-Entity dataset uses instance segmentation and repainting to acquire a substantial volume of training data.
- Dataset construction: The web-images dataset uses cross-verification with an MLLM to retain group photos whose individual identities are confirmed.For group photos, sampled individual images are checked against the group image.