UniGP: Taming Diffusion Transformer for Prior-Preserved Unified Generation and Perception

2026-06-29

AI Summary

Existing approaches treat controllable generation with diffusion models and dense prediction as disjoint tasks, overlooking the potential of jointly modeling their heterogeneous output distributions. This work proposes the Dense Unified Generation and Perception (DUGP) module and a unified dataset training strategy built upon the MMDiT architecture, enabling non-RGB dense outputs to be modeled simply by duplicating the image branch—thereby achieving mutual enhancement between generative and perceptual tasks. Adhering to Occam’s razor, the method requires no task-specific designs and instead learns the joint distribution of image-geometry pairs through multi-task co-training. The resulting unified model surpasses prior unified approaches and matches the performance of specialized models: generative priors enrich perceptual detail, while perception signals improve structural alignment in generation.

Abstract

Recent advances in diffusion models have shown impressive performance in controllable image generation and dense prediction tasks. However, existing approaches typically treat diffusion-based controllable generation and dense prediction as separate tasks, overlooking the potential benefits of jointly modeling the heterogeneous distributions. In this work, we introduce UniGP, a framework built upon MMDiT, which unifies controllable generation and dense prediction through simple joint training, without the need for complex task-specific designs or losses, while preserving the backbone's versatile priors. By learning controllable generation and prediction under different conditions, our model effectively captures the joint distribution of image-geometry pairs. UniGP is capable of versatile controllable generation, dense prediction, and joint generation. Specifically, the proposed UniGP consists of DUGP and a unified dataset training strategy. The former, following the principle of Occam's razor, uses only a copied image branch of MMDiT to model dense distributions beyond RGB, while the latter integrates heterogeneous datasets into a unified training framework to jointly model generation and perception tasks. Extensive experiments demonstrate that our unified model surpasses prior unified approaches and performs on par with specialized methods. Furthermore, we demonstrate that multi-task joint training provides complementary benefits: generative priors enrich perceptual details, while perceptual learning improves structural alignment in generation.

Problem

• diffusion models
• controllable generation
• dense prediction
• unified modeling
• heterogeneous distributions

Innovation

• diffusion transformer
• unified generation and perception
• joint training
• dense prediction
• controllable generation