OmniGen: Unified Image Generation

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.

한 줄 요약

OmniGen은 단일 모델로 텍스트-이미지 생성, 이미지 편집, 조건 기반 생성 등 다양한 이미지 생성 작업을 처리하는 통합 확산 모델이다.

핵심 기여도

**OmniGen**이라는 통합 이미지 생성 모델을 제안, 텍스트-이미지 생성 외에도 이미지 편집, 주체 기반 생성, 시각 조건 생성을 지원.
**X2I**라는 대규모 다중 작업 이미지 생성 데이터셋을 구축, 모든 작업을 통일된 포맷으로 표준화.
**단일 VAE + 트랜스포머** 구조로 기존 확산 모델의 복잡한 플러그인 없이 작업을 단계 없이 처리.
**이미지 편집 작업에서 가중치 손실(weighted loss)** 을 도입해 모델이 입력 이미지를 그대로 복사하는 단축 학습을 방지.

핵심 아이디어

기존 이미지 생성 모델은 각 작업에 맞는 별도의 모듈(예: ControlNet, InstructPix2Pix)이 필요했으며, 복잡한 워크플로우를 요구했다. OmniGen은 대규모 언어 모델(LLM)이 언어 생성을 통합하는 방식을 이미지 생성에 적용한 첫 시도이다. 모델은 **단일 VAE와 트랜스포머**로 구성되어 있으며, 텍스트와 이미지를 자유롭게 결합한 입력을 처리할 수 있다. 특히, **이미지 패치 간 양방향 어텐션**을 도입해 이미지 전체를 통합적으로 모델링하고, **KV 캐시 기반 추론 가속**을 통해 LLM과 유사한 효율성을 달성했다. 이는 기존 확산 모델에서 별도 인코더(예: CLIP)를 사용하는 방식과 대조된다.

기술적 접근법

**모델 아키텍처**: VAE와 트랜스포머로 구성. VAE는 SDXL에서 가져오고, 트랜스포머는 Phi-3로 초기화.
**입력 표현**: 텍스트는 Phi-3 토크나이저로 처리, 이미지는 VAE로 잠재 표현 추출 후 2×2 패치로 분할.
**어텐션 메커니즘**: 텍스트는 인과적 어텐션, 이미지는 양방향 어텐션 적용.
**추론**: 가우시안 노이즈 샘플링 후, **flow matching**을 사용해 여러 단계 반복.
**데이터셋**: X2I 데이터셋을 사용, 다양한 작업을 통일된 포맷으로 표준화.
**가중치 손실**: 이미지 편집 작업에서 작은 영역만 수정되므로, 가중치 손실을 통해 단축 학습 방지.

주요 결과

**X2I 데이터셋**에서 OmniGen은 텍스트-이미지 생성, 이미지 편집, 조건 기반 생성 등 다양한 작업에서 뛰어난 성능을 보임.
**가중치 손실 미사용 시**, 이미지 편집 작업에서 입력 이미지를 그대로 복사하는 문제가 발생.
**단일 VAE + 트랜스포머** 구조로 기존 확산 모델 대비 워크플로우 단계를 줄이고, 추론 속도를 향상.
**KV 캐시 기반 추론**으로 LLM과 유사한 추론 효율성 달성.

의의 및 한계

OmniGen은 이미지 생성 분야에서 LLM이 언어 생성을 통합한 것처럼, 다양한 작업을 단일 모델로 처리하는 기초 모델의 첫 사례이다. 특히, **X2I 데이터셋**을 통해 다양한 작업을 통일된 포맷으로 학습할 수 있으며, **단일 VAE + 트랜스포머** 구조로 모델의 복잡도를 줄였다. 그러나, **모델의 추론 정확도나 생성 품질에 대한 정량적 수치**는 명시되지 않았으며, **대규모 실제 이미지 생성 환경에서의 성능**도 추가 연구가 필요하다. 또한, **모델의 추론 속도 향상 폭**(예: 1.6× 가속 등)도 명시되지 않았다.

실용적 활용

OmniGen은 디자인, 콘텐츠 제작, 게임 개발 등에서 사용자 지시에 따라 다양한 이미지 생성 작업을 단계 없이 처리할 수 있어, **복잡한 워크플로우를 대체하는 실용적 도구**로 활용 가능하다. 또한, **단일 모델로 다양한 작업을 처리**할 수 있어, 기업이나 개발자들이 별도의 모델을 유지하는 부담을 줄일 수 있다.