Source-linked AI summary
UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild
Can Qin, Shu Zhang, Ning Yu, Yihao Feng, Xinyi Yang, Yingbo Zhou, Huan Wang, Juan Carlos Niebles, Caiming Xiong, Silvio Savarese, Stefano Ermon, Yun Fu, Ran Xu
TL;DR
Existing visual generative models provide flexible language control but do not adequately unify precise spatial, structural, and geometric controls across visual conditions. UniControl augments text-to-image diffusion models with an MOE-style adapter and task-aware HyperNet, and achieves strong multi-task performance with zero-shot generalization to unseen tasks. Its main scope boundary is dependence on biased training data and the broader limitations of diffusion-based image generation.
Problem
Visual generative models lack a unified framework that combines arbitrary language prompts with pixel-level spatial, structural, and geometric control across diverse visual conditions.
Method
UniControl augments pretrained text-to-image diffusion models with an MOE-style adapter and task-aware HyperNet to handle diverse controllable generation tasks simultaneously.
Results
UniControl outperforms single-task controlled image-generation methods and demonstrates zero-shot generation on unseen tasks.
Takeaways & Limitations
A single model can provide versatile controllable generation across multiple visual conditions while retaining language conditioning.
Takeaways & Limitations
UniControl remains limited by bias in its Laion-Aesthetics-derived training data and may produce biased or low-fidelity outputs.
Abstract
from arXiv · showhide
Achieving machine autonomy and human control often represent divergent objectives in the design of interactive AI systems. Visual generative foundation models such as Stable Diffusion show promise in navigating these goals, especially when prompted with arbitrary languages. However, they often fall short in generating images with spatial, structural, or geometric controls. The integration of such controls, which can accommodate various visual conditions in a single unified model, remains an unaddressed challenge. In response, we introduce UniControl, a new generative foundation model that consolidates a wide array of controllable condition-to-image (C2I) tasks within a singular framework, while still allowing for arbitrary language prompts. UniControl enables pixel-level-precise image generation, where visual conditions primarily influence the generated structures and language prompts guide the style and context. To equip UniControl with the capacity to handle diverse visual conditions, we augment pretrained text-to-image diffusion models and introduce a task-aware HyperNet to modulate the diffusion models, enabling the adaptation to different C2I tasks simultaneously. Trained on nine unique C2I tasks, UniControl demonstrates impressive zero-shot generation abilities with unseen visual conditions. Experimental results show that UniControl often surpasses the performance of single-task-controlled methods of comparable model sizes. This control versatility positions UniControl as a significant advancement in the realm of controllable visual generation.
1 Introduction
UniControl addresses the difficulty of combining precise visual control with flexible language conditioning in a single vision-generation model. It unifies nine controllable generation tasks using task-aware adaptation and demonstrates improved multi-task and zero-shot capabilities.
- Stable Diffusion provides flexible text control but lacks pixel-level precision for spatial, structural, and geometric image controls.
- UniControl combines language conditioning with diverse visual conditions in one unified controllable diffusion model.Its design seeks a universal representation space across tasks while preserving visual control and textual guidance.
- The model uses an MOE-style adapter for task-specific low-level features and a task-aware HyperNet that modulates ControlNet using task instructions.The HyperNet converts task instructions into embeddings that modulate ControlNet’s zero-convolution layers.
- UniControl is pretrained on nine tasks across five categories using the MultiGen-20M dataset of more than 20 million image-text-condition triplets.
- UniControl outperforms each single-task controlled image-generation method and demonstrates zero-shot generalization to unseen tasks.
2 Related Works
Related work establishes diffusion models as strong foundations for text-to-image generation, image translation, and editing. UniControl builds on these diffusion-based lines of research to add unified visual-condition control.
- Diffusion models use training and sampling advances to generate images effectively, including score-based diffusion, DDPM, and DDIM.
- Text-to-image diffusion models condition U-Net denoisers on pretrained text embeddings, supporting image generation and editing.
- ControlNet and UniControl build on diffusion-based text-to-image generation, while compositional conditions provide another route for guiding visual generation.
- Image-to-image translation methods learn mappings between image domains, while newer diffusion approaches support translation, editing, colorization, superresolution, and inpainting.
3 UniControl
UniControl unifies diverse visual-condition controls and language prompts in one diffusion model. Its MOE-style adapter and task-aware HyperNet support task-specific conditioning and zero-shot generalization to hybrid and unseen tasks.
- 3.1 Training Setup: UniControl reformulates training pairs with task instructions, language prompts, visual conditions, and target images so one model can distinguish multiple tasks.The training objective samples tasks uniformly and optimizes the corresponding diffusion loss.
- 3.2 Model Design: The model is designed to address low-level feature misalignment while learning shared knowledge and task differences across visual conditions.This supports conditions such as segmentation maps and depth maps without requiring separate models for each task.
- 3.2 Model Design: Figure 2 combines the MOE adapter and HyperNet in a singular multi-task model, reducing size relative to an equivalent stack of single-task models.The adapter uses roughly 70K parameters per task, while the HyperNet has roughly 12M parameters and modulates seven zero-convolution layers.
- 3.2 Model Design: The MOE-style adapter uses task-specific convolution modules to capture unique low-level features and retrieve related adapters for zero-shot tasks.Each convolution module has approximately 70K parameters and explicitly differentiates visual conditions.
- 3.3 Task Generalization Ability: UniControl demonstrates zero-shot generalization to hybrid conditions and unseen visual tasks through convolution-output integration and estimated combinations of pretrained task adapters.Hybrid examples combine segmentation maps with human skeletons; unseen-task weights may be manually assigned or estimated from task-instruction similarity.
- 3.2 Model Design: The task-aware HyperNet converts task instructions into embeddings and uses them to modulate ControlNet zero-convolution layers for task-aware visual control.The embedding length matches each layer’s input channels, with elements scaling convolution kernels per input channel.
4 Experiments
UniControl is evaluated as a unified multi-task model across diverse visual conditions, against task-specific baselines and single-task variants. Experiments examine visual quality, user preferences, FID, ablations, hybrid conditions, and zero-shot generalization to unseen tasks.
- Experiment setup: Experiments evaluate UniControl across diverse conditions and tasks using comprehensive empirical studies.The model is assessed with visual comparisons, user studies, FID analysis, ablations, hybrid conditions, and unseen-task generalization.
- Implementation: UniControl is implemented with Stable Diffusion, ControlNet, an MOE adapter, and a task-aware HyperNet in approximately 1.5B parameters.The HyperNet uses task-instruction embeddings to modulate zero-convolution kernels.
- Visual comparison: Visual comparisons across eight tasks show UniControl consistently outperforming ControlNet in visual quality and alignment with conditions or prompts.Reported improvements include detail preservation for Canny, sharper boundaries for HED, and better condition alignment for segmentation, OpenPose, bounding boxes, and outpainting.
- User study: User studies report that UniControl outperforms official ControlNet checkpoints on all six evaluated tasks.The advantage is especially pronounced for HED, depth, and surface-normal generation, where users favor UniControl.
- Quantitative evaluation: FID evaluations on over 2,000 Laion and COCO test samples show UniControl consistently surpassing GLIGEN, T2I-Adapter, and other baselines across multiple task categories.The evaluated categories include edges, regions, skeletons, and geometric maps.
- Ablation study: The full UniControl model significantly outperforms ablations lacking the MOE-style adapter and task-aware HyperNet.This ablation compares FID scores for the full model against component-removal variants.
- Zero-shot generalization: UniControl generates results for unseen hybrid condition combinations and unseen deblurring, colorization, and inpainting tasks.The reported examples attribute this zero-shot behavior to shared attributes and implicit correlations among pre-training and new tasks.
5 Conclusion and Discussion
UniControl is presented as a unified model for incorporating diverse conditions into diffusion-based generation. The authors report robust performance and adaptability, while noting limitations from biased training data and the broader constraints of diffusion models.
- Conclusion: UniControl incorporates a wide range of conditions into diffusion-model generation through a unified design.Its two key components are a Mixture-of-Experts-style adapter and a task-aware HyperNet.
- Conclusion: The experiments report robust performance and adaptability across different tasks and conditions.The stated potential concerns complex text-to-image generation tasks.
- Limitation and Broader Impact: UniControl is limited by bias in its subset of Laion-Aesthetics training data and may produce biased or low-fidelity outputs.The authors also note limitations when high-quality human output is desired.
A Details of Implementation
The MOE adapter uses parallel task-specific convolutional modules while sharing the remaining model parameters across tasks.
- Details of Implementation: The MOE adapter comprises nine parallel ConvNets, each with three convolutional and nonlinear-activation layers and 70K parameters.Task keys designate the adapter corresponding to each visual condition.
- Details of Implementation: After task-specific MOE processing, the remaining model parameters are shared across all tasks.This architecture supports task adaptability while concentrating task-specific processing in the adapters.
A.2 Task-aware HyperNet
The task-aware HyperNet modulates ControlNet using embeddings derived from natural-language task instructions.
- Task-Aware HyperNet: The task-aware HyperNet modulates ControlNet’s zero-convolution modules using the task instruction condition.It projects task instructions into task embeddings with a CLIP text encoder and injects them into zero-convolution layers.
- Task-Aware HyperNet: The HyperNet implements task-aware control by multiplying task embeddings into the zero-convolution layers of a trainable ControlNet copy.The embedding length matches the number of input channels in each zero-convolution layer.
A.3 Data Collection
MultiGen-20M was assembled by extracting nine visual conditions from filtered source images, yielding over 20 million condition-image-prompt triplets.
- Data Collection: MultiGen-20M contains over 20 million condition-image-prompt triplets spanning nine visual-condition tasks.The collection uses original images, language prompts, and visual conditions across the task set.
B Numerical Analysis of Task-Aware Modulated ControlNet
The task-aware modulated ControlNet preserves zero initialization while remaining trainable after optimization begins.
- Zero Initialization: The modified ControlNet retains ControlNet’s zero-initialization property for its convolution layers.Weights and biases of the convolution layers are initialized as zeros.
- Model Formulation: The output of task-aware modulated ControlNet combines the frozen Stable Diffusion block with condition processing modulated by task instructions.The formulation defines the input feature map, visual condition, task instruction, zero convolutions, and task-aware HyperNet.
- Gradient Analysis: After the first update, the zero-convolution output becomes nonzero, allowing the network to follow standard stochastic-gradient-descent learning dynamics.The analysis states that the task-aware HyperNet can then receive nonzero gradients and that the model can be efficiently updated.
- Gradient Analysis: Before the first gradient update, zero convolution output is zero, so the task-aware HyperNet receives zero gradient through that path.This follows from the zero-initialized convolution output and the gradient analysis of the modulated component.
C Zero-shot-task Results and Analysis
UniControl shows promising zero-shot generation on untrained tasks, including deblurring, colorization, inpainting, and scribble-guided generation, while retaining visible failure modes.
- Zero-shot Results: UniControl produces zero-shot deblurring, colorization, and inpainting results despite not being trained on those tasks.The authors report recovered fine-grained details, generally acceptable colorization, and impressive inpainting, with some missing details and imperfect faces.
- Inpainting versus Outpainting: Outpainting weights transfer imperfectly to inpainting because inpainting requires contextual consistency around masked regions, unlike outpainting’s freer content generation.Direct reuse can create sharp changes at mask boundaries.
- Scribble Generalization: UniControl generalizes to scribble conditions even though scribble data were unseen during training, showing parallels to ControlNet-Scribble.The comparison is illustrated against a model trained directly on scribble data.
- User Study: The evaluation used three Amazon Mechanical Turk master workers to select the better output between paired generated images.Annotators received the condition map and descriptive text before choosing between two outputs.
- Failure Cases: The reported failure cases include distorted bodies, blurred faces, and unrealistic images, which the authors associate with Stable Diffusion bias and training-data limitations.The authors suggest improving the training data quantitatively and qualitatively.
F Additional Results
Additional visualizations cover the trained condition types and show realistic outputs across edges, geometry, pose, regions, editing, and user sketches.
- Additional Results: Additional results visualize Canny, HED, depth, surface normal, human skeleton, bounding-box, segmentation, and outpainting generation.These figures provide qualitative examples across the principal trained tasks.
- User Sketch: UniControl generates promising realistic images from user sketches, with the corresponding user-study results reported separately.The user study is described as showing more votes for UniControl than the single-task model.
- Failure Cases: A failure-case figure documents distorted bodies, blurred faces, and incorrect image creation among the generated outputs.These examples are presented as qualitative failure cases.
- Zero-shot Editing: The additional figures also include zero-shot deblurring, gray-to-RGB colorization, and image in-painting examples.The in-painting adapter weights are directly inherited from outpainting.