PaperlayerView exact revision on arXiv

Source-linked AI summary

UniVG: A Generalist Diffusion Model for Unified Image Generation and Editing

Tsu-Jui Fu, Yusu Qian, Chen Chen, Wenze Hu, Zhe Gan, Yinfei Yang

arXiv:2503.12652v2cs.CV

TL;DR

Task-specific diffusion models require separate architectures, training designs, and parameter sets. UniVG unifies diverse image-generation tasks in one generalist diffusion model and studies data mixing and multi-task training. It handles the considered tasks without compromising core text-to-image performance, with a GenEval score of 0.70 versus FLUX.1-dev's 0.66.

  • Problem

    Task-specific diffusion systems require separate model architectures, training designs, and parameter sets, creating management and computational-resource challenges.

    Source: §1 · p. 2

  • Method

    UniVG uses a unified diffusion framework that adapts to diverse tasks through multimodal inputs and studies modeling, data curation, and training strategies.

    Source: §1 · p. 2, §3.2 · p. 4, §3.3 · p. 4

  • Results

    UniVG handles the considered image-generation tasks without compromising core text-to-image performance; its GenEval score is 0.70 versus FLUX.1-dev's 0.66.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    The results support a single model effectively handling the considered tasks without sacrificing core text-to-image performance.

    Source: §5 · p. 8

Abstract

from arXiv · show

Text-to-Image (T2I) diffusion models have shown impressive results in generating visually compelling images following user prompts. Building on this, various methods further fine-tune the pre-trained T2I model for specific tasks. However, this requires separate model architectures, training designs, and multiple parameter sets to handle different tasks. In this paper, we introduce UniVG, a generalist diffusion model capable of supporting a diverse range of image generation tasks with a single set of weights. UniVG treats multi-modal inputs as unified conditions to enable various downstream applications, ranging from T2I generation, inpainting, instruction-based editing, identity-preserving generation, and layout-guided generation, to depth estimation and referring segmentation. Through comprehensive empirical studies on data mixing and multi-task training, we provide detailed insights into the training processes and decisions that inform our final designs. For example, we show that T2I generation and other tasks, such as instruction-based editing, can coexist without performance trade-offs, while auxiliary tasks like depth estimation and referring segmentation enhance image editing. Notably, our model can even outperform some task-specific models on their respective benchmarks, marking a significant step towards a unified image generation model.

1. Introduction

UniVG targets the management burden of task-specific diffusion systems by unifying diverse image-generation tasks in one model. The paper studies modeling, data recipe, and training strategy, reporting that instruction editing coexists with text-to-image generation without compromise.

1. Introduction

  • The work addresses the management and computational burden of task-specific models, while examining modeling, data recipe, and training choices left unclear by prior generalist-model studies.The authors identify a lack of clear ablations on optimal training design in OmniGen and OneDiff.Source: §1 · p. 2
  • UniVG integrates text prompts, masks, and images in a minimally modified MM-DiT, with channel-wise latent concatenation improving efficiency over OmniGen.External conditions such as semantic maps or user-defined attributes can be injected through embedding replacement.Source: §1 · p. 2
  • Instruction-based editing does not compromise text-to-image generation, while depth estimation and referring segmentation enhance image editing.These findings come from the paper's investigation of synergy among tasks.Source: §1 · p. 2
  • Training proceeds progressively: large-scale text-to-image pretraining, gradual introduction of editing and other tasks, then ID-preserving generation fine-tuning.The final stage mixes in additional ID-preserving generation data.Source: §1 · p. 2
  • UniVG reports state-of-the-art performance against OmniGen and OneDiff, with a GenEval score of 0.70 versus FLUX.1-dev's 0.66.The paper also reports strong performance across all image-generation tasks considered.Source: §1 · p. 2

2. Related Work

Related work traces diffusion progress from latent-space models toward flow-based formulations and transformer architectures, alongside specialized control, editing, and personalization methods. Recent generalist models unify image-generation tasks; UniVG presents a more thoroughly studied training recipe.

2. Related Work

  • Latent-space diffusion work increasingly adopts flow-based formulations and transformer architectures, with SD3 and Flux.1 presented as state-of-the-art open-source models.The passage says flow-based approaches simplify generation and improve convergence speed and generation quality.Source: §2 · p. 2
  • Prior work explores fine-grained control, instruction-based editing, and personalized generation, while other approaches unify multiple image conditions or image understanding and generation.The passage notes that some multi-condition methods require complex adapters for each condition.Source: §2 · pp. 2–3
  • OmniGen and OneDiff are cited as generalist models for diverse image-generation tasks, while UniVG emphasizes a more thoroughly studied training recipe.UniVG's overview also describes support for text-to-image, in/outpainting, instruction editing, layout guidance, identity preservation, and auxiliary perception tasks.Source: §2 · p. 3, Figure 2 · p. 3

3. Method

UniVG combines multimodal inputs in one diffusion framework and adapts to different image-generation tasks through task-specific input formats. Its method uses flow matching, external-condition injection, and staged multi-task training.

3.1. Background

  • MM-DiT fuses modalities through structured attention, while conditional flow matching learns a continuous noise-to-data transformation that supports ODE sampling with fewer steps.The model learns to imitate the ideal transformation because the time-dependent velocity field is generally intractable.Source: §3.1 · p. 3

3.2. UniVG

  • UniVG interpolates noise with the encoded target image, then concatenates noisy latents and visual inputs along channels for efficient flow-matching training.The channel-wise design is contrasted with sequence-wise concatenation and is described as mitigating context-perception disruption.Source: §3.2 · pp. 3–4
  • External conditions are encoded into features that replace placeholder-token embeddings, allowing additional guidance in formats beyond images.The feature length can vary if multiple placeholder tokens are available.Source: §3.2 · p. 4

3.3. Multi-task Training

  • Text-to-image and in/outpainting use task tokens and masks, while instruction editing uses the source image and instruction with all regions editable.For background in/outpainting, the prompt is discarded; the complete image is the target output.Source: §3.3 · p. 4
  • Auxiliary depth estimation, pose estimation, and referring segmentation train on visualized outputs; layout guidance uses a rendered layout, while identity customization injects facial features.The auxiliary-task outputs are treated as images, and identity features are extracted with a CLIP encoder followed by a two-layer MLP.Source: §3.3 · p. 4
  • Training proceeds from text-to-image foundation training to multi-task training, then fine-tunes identity customization alongside other tasks after Stage II showed catastrophic forgetting.Stage III uses a 1:1 ratio of identity-preserving data to all other multi-task data.Source: §3.3 · pp. 4–5

4. Experiments

The experiments evaluate UniVG across text-to-image generation, editing, identity-preserving generation, and auxiliary tasks, alongside studies of training, scaling, and inference efficiency. Results and ablations examine task performance and interactions within the unified model.

4.1. Experimental Setup

  • The dataset collection covers T2I, in/outpainting, instruction editing, depth and pose estimation, referring segmentation, layout guidance, and ID-preserving generation.Stage III uses a 1:1 ratio of ID-preserving generation to all other tasks.Source: §4.1 · p. 5

4.2. Evaluation Results

  • UniVG surpasses unified baselines on text-to-image compositionality and semantics, remains competitive with SD3, and reports higher editing CLIP-T than task-specific methods.For instruction-based editing, UniVG also has advantages over OmniGen in both CLIP-T and CLIP-I.Source: §4.2 · pp. 5–6, Table 3 · p. 5
  • On Unsplash-50, UniVG outperforms unified baselines in face consistency and achieves higher CLIP-T for identity-preserving generation.The authors report that all-at-once training gives worse ID-preservation performance.Source: §4.2 · p. 6, Table 4 · p. 6, Table 5 · p. 6

4.3. Ablation Study

  • T2I and editing coexist, but ID training drops MagicBrush CLIP-T from 28.3 to 26.9; auxiliary tasks improve editing, and the full recipe balances performance.Table 6 compares stages that add instruction editing, auxiliary tasks, layout-guided generation, and ID-preserving generation.Source: §4.3 · pp. 6–7, Table 6 · p. 7
  • Scaling B (416M), L (1.8B), and XL (3.7B) improves performance and generalizes across generation and editing; layout prompts guide object positions and counting.Figure 5 presents the scaling results; Figure 6 shows layout-guided examples.Source: Figure 5 · p. 7, §4.3 · p. 8, Figure 6 · p. 8
  • Channel-wise concatenation preserves sequence length across generation and editing; UniVG is more efficient than OneDiff for editing, while OmniGen requires 36+ seconds.The efficiency comparison uses 512 resolution and reports inference time and GPU cost.Source: §4.3 · p. 7, Table 7 · p. 7
  • UniVG completes or expands missing image regions with or without text guidance, and also supports depth and pose estimation and referring segmentation.The examples include text-guided insertion of specified text or objects and segmentation of objects or backgrounds.Source: §4.3 · p. 8, Figure 7 · p. 8, Figure 8 · p. 8

5. Conclusion

UniVG unifies diverse image-generation tasks in one diffusion model, using a minimalist architecture and a three-stage training pipeline. The authors report that it handles the tasks without compromising core text-to-image performance.

5. Conclusion

  • A single UniVG model effectively handles diverse image-generation tasks without compromising core text-to-image performance.The conclusion describes this as the result of a generalist model within one framework.Source: §5 · p. 8
  • The final design adopts a minimalist architecture and a three-stage training pipeline after ablations of key modeling and data choices.Source: §5 · p. 8
  • Future work will add vision perception tasks and scale up the model.Source: §5 · p. 8

Related papers

Loading 2503.12652v2…