Title: UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild
Venue: NeurIPS 2023
Accepted: yes
Keywords: Image Generation

Multi-modal

HyperNet
Abstract: 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.

UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild

Abstract

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

1 Introduction

Generative foundation models are revolutionizing the ways that humans and AI interact in natural language processing (NLP) [1-6], computer vision (CV) [7-10], audio processing (AP) [11, 12], and robotic controls [13-15], to name a few. In NLP, generative foundation models such as InstructGPT or GPT-4, achieve excellent performance on a wide range of tasks, e.g., question answering, summarization, text generation, or machine translation within a single-unified model. Such multi-tasking ability is one of the most appealing characteristics of generative foundation models. Furthermore, generative foundation models can also perform zero-shot or few-shot learning on unseen tasks [3, 16, 17].

For generative models in vision domains [9, 18-20], such multi-tasking ability is less clear. Stable Diffusion Model (SDM) [9] has established itself as the major cornerstone for text-conditioned image generation. However, while text descriptions provide a very flexible way to control the generated images, their ability to provide pixel-level precision for spatial, structural, or geometric controls is

'Chocolate mousee cake'

HED-to-Image

'Three zebras grazing in a grassy area near shrubs'

Depth-to-Image

'One tray of rice and a tray of fruits and veggies.'

'Awesome 3d wall panels and interior wall paneling ideas for Living room 3d tiles'

Seg-to-Image

'Casual Short Hairstyles For Women'

Normal-to-Image

'A man with glasses and wearing a tie poses for a profile picture'

BBox-to-Image

'A man in a Rays uniform is throwing a baseball '

Skeleton-to-Image

'Breath of the Wild compressed cover art'

Canny-to-Image

Pre-training Tasks

Input Condition

Ours Result

Ours Result

UniControl Diffusion Model

to Image

HED Edge to

Image

Bounding Box

to Image

Segmentation

to Image

Normal Map

to Image

to Image

Skeleton to

Image

Image

Outpainting

Skeleton to Image

Segmentation +

Skeleton to Image

Gray Image Colorization

Image Deblur

Image Inpainting

…

…

…

Pre-training Tasks

Zero-shot Tasks

Zero-shot Combination

Zero-shot Task: Image Deblurring

'Arches National Park, Three Gossips, With Milky Way In The Background'

Zero-shot Task: Depth-and-Skeleton-to-Image

'Background: White House, Foreground: A group of women sitting on a bench'

Zero-shot Task: Image Colorization

'A cave on a desert planet'

Zero-shot Task: Seg-and-Skeleton-to-Image

'Background: Old City of Tzfat, Foreground: A set of twins who are taking a bath'

often inadequate. A recent work, ControlNet [21], was proposed to augment SDM to enable visual conditions ( e.g., edge maps, depth maps). With the additional visual conditions, ControlNet can achieve explicit spatial, structural, or geometric control over generated structures, without losing the semantic control from textual captions. Unfortunately, unlike language prompts that a unified module such as CLIP [22] can handle, each ControlNet model can only handle a specific control modality that it was trained on ( e.g. , edge map). Retraining a separate model is necessary to handle a different modality of visual conditions, incurring non-trivial time and spatial complexity costs.

To overcome the limitation of previous works, we present UniControl , a unified diffusion model for controllable visual generation in the wild, which is capable of simultaneously handling both

language and various visual conditions. Naturally, UniControl can perform multi-tasking and can encode visual conditions from different tasks into a universal representation space, seeking a common representation structure among tasks. The unified design of UniControl allows us to enjoy the advantages of improved training and inference efficiency, as well as enhanced controllable generation. On the one hand, the model size of UniControl does not significantly increase as the number of tasks scales up. On the other hand, UniControl derives advantages from the inherent connections between different visual conditions [ e.g. , 23-25]. These relationships, such as depth and segmentation mapping, leverage shared geometric information to enhance the controllable generation quality.

The unified controllable generation ability of UniControl relies on two novel designed modules, a mixture of expert (MOE)-style adapter and a task-aware HyperNet [26, 27]. The MOE-style adapter can learn necessary low-level feature maps from various visual conditions, allowing UniControl to capture unique information from different visual conditions. The task-aware HyperNet, which takes the task instruction as natural language prompt inputs, and outputs a task-aware embedding. The output embeddings can be incorporated to modulate ControlNet [21] for task-aware visual condition controls, where each task corresponds to a particular format of visual condition. As a result, the task-aware HyperNet allows UniControl to learn meta-knowledge across various tasks, and obtain abilities to generalize to unseen tasks. As Tab. 1, UniControl has significantly compressed the model size compared with its direct baseline, i.e. , Multi-ControlNet, by unifying nine tasks into ONE model.

To obtain multi-tasking and zero-shot learning abilities, we pre-train UniControl on nine distinct tasks across five categories: 1) edges (Canny, HED, User Sketch); 2) region-wise maps (Segmentation Maps, Bounding Boxes); 3) skeletons (Human Pose Skeletons); 4) geometric maps Depth, Surface Normal); 5) editing (Image Outpainting). We build MultiGen-20M dataset, comprising over 20 million high-quality triplets of original images, language prompts, and visual conditions for all the tasks. Then UniControl is trained for over 5,000 GPU hours on NVIDIA A100-40G hardware that is comparable with the overall training cost of different ControlNets. Moreover, UniControl exhibits a remarkable capacity for zero-shot adaptation to new tasks, highlighting its potential for deployment in real-world applications. Our contributions are summarized below:

We present UniControl, a unified model capable of handling various visual conditions for the controllable visual generation.

We collect a new dataset for multi-condition visual generation with more than 20 million imagetext-condition triplets over nine distinct tasks across five categories.

We conduct extensive experiments to demonstrate that the unified model UniControl outperforms each single-task controlled image generation, thanks to learning the intrinsic relationships between different visual conditions.

UniControl shows the ability to adapt to unseen tasks in a zero-shot manner, highlighting its versatility and potential for widespread adoption in the wild.

2 Related Works

Diffusion-based Generative Models. Diffusion models were initially introduced in [28] that yield favorable outcomes for generating images [18, 21]. Improvements have been made through various training and sampling techniques such as score-based diffusion [29, 30], Denoising Diffusion Probabilistic Model (DDPM) [31], and Denoising Diffusion Implicit Model (DDIM) [32], When training U-Net denoisers [33] with high-resolution images, researchers involve speed-up techniques including pyramids [34], multiple stages [20], or latent representations [9]. In particular, UniControl leverages Stable Diffusion Models (SDM) [9] as the base model to perform multi-tasking.

Text-to-Image Diffusion. Diffusion models emerge to set up a cutting-edge performance in text-to-image generation tasks [20, 19], by cross-attending U-Net denoiser in diffusion generators with CLIP [22] or T5-pretrained [2] text embeddings. GLIDE [35] is another example of a textguided diffusion model that supports image generation and editing. UniControl and closely related

ControlNet [21] are both built upon previous works on diffusion-based text-to-image generation [9]. [36] introduces the compositional conditions to guide visual generation.

Image-to-Image Translation. Image-to-image (I2I) translation task was initially proposed in Pix2Pix [37], focusing on learning a mapping between images in different domains. Recently, diffusion-based approaches [38, 39, 21] set up the new state of the art results. Recent diffusionbased image editing methods show outstanding performances without requiring paired data, e.g., SDEdit [40], prompt-to-prompt [41], Edict [42]. Other image editing examples include various diffusion bridges and flows [43-47], classifier guidance [30] based methods for colorization, superresolution [34], inpainting [48], and etc . ControlNet [21] takes both visual and text conditions and achieves new state-of-the-art controllable image generation. Our proposed UniControl unifies various visual conditions of ControlNet, and is capable of performing zero-shot learning on newly unseen tasks. Concurrently, Prompt Diffusion [49] introduces visual prompt [50] from image inpainting to controllable diffusion models, which requires two additional image pairs as the in-context example for both training and inference. By contrast, UniControl takes only a single visual condition while still capable of both multi-tasking and zero-shot learning.

3 UniControl

In this section, we describe the training and the model design of our unified controllable diffusion model UniControl . Specifically, we first provide the problem setup and training objectives in Sec. 3.1, and then show the novel network design of UniControl in Sec. 3.2. Finally, we explain how to perform zero-shot image generation with the trained UniControl in Sec. 3.3.

3.1 Training Setup

Different from the previous generative models such as Stable Diffusion Models (SDM) [9] or ControlNet [21], where the image generation conditions are single language prompt, or single type of visual condition such as canny , UniControl is required to take a wide range of visual conditions from different tasks, as well as the language prompt.

To achieve this, we reformulate the training conditions and target pairs for UniControl. Specifically, suppose we have a dataset consisting of K tasks : D := {D 1 ∪··· ∪ D K }, and for each task training set D k , denote the training pairs by ([ c text , c task ] , I c , x ) , with c task being the task instruction that indicates the task type, c text being the language prompt describing the target image, I c being the visual conditions, and x being the target image. With the additional task instruction, UniControl can differentiate visual conditions from different tasks. A concrete training example pair is the following:

Task-Aware Vision-Language Condition

'Camp on a mountain top: Birthday Presents, Adventure, Outdoor, Mountain Camps, Great View, Places, Hiking, Mornings Lights, Himalayan Sunri'

'Canny Edge to Image' Task Instruction               :

Visual Condition        :

Language Prompt               :

Target output

where the task is to translate the canny edge to real images following language prompt. With the induced training pairs ( x, [ c task , c text ] , I c ) , we define the training loss for task k following LDM [9]:

ℓ k ( θ ) := E z,ε,t,c task ,c text , I c [ ∥ ε -ε θ ( z t , t, c task , c text , I c ) ∥ 2 2 ] , with ([ c task , c text ] , I c , x ) ∼ D k , where t represents the time step, z t is the noise-corrupted latent tensor at time step t , z 0 = E ( x ) , and θ is the trainable parameters of UniControl . We also apply classifier-free guidance [51] to randomly drop 30% text prompts to enhance the controllability of input visual conditions. We train UniControl uniformly on the K tasks. To be more specific, we first randomly select a task k and sample a mini-match from D k , and optimize θ with the calculated loss ℓ k ( θ ) .

3.2 Model Design

Since our unified model UniControl needs to achieve superior performance on a set of diverse tasks, it is necessary to ensure the network design enjoys the following properties: 1) The model can overcome

UniControl Model Overview

SD Middle

(Trainable

Copy)

Modulated Zero Conv

Modulated Zero Conv

SD Encoder

(Trainable

Copy)

Task

Instruction

Text &

Time

Modulated Zero Conv

SD Middle

SD Decoder

SD Block

Trainable

Copy of

SD Block

x

Condition

*

Zero Convolution

Task Embedding

*

Zero Convolution

Task Embedding

y

Modulated Zero Conv

Task-Aware Vision-Language Condition

Input one or more visual conditions

* N

* N

Generated

Images

'a young child looking at

a birthday cupcake'

SD Encoder

Feature

Conv

Con

v

Con

v

Conv

...

…

…

Task Instruction:

"normal

surface to image"

Task Embedding

Task-Aware

HyperNet

Trainable

Weights

CLIP Text Encoder

Base Hypernet

Parallel Linear Layers

Linear Layer

Task-Aware

HyperNet

1

the misalignment of low-level features from different tasks; 2) The model can learn meta-knowledge across tasks, and adapt to each task effectively.

The first property can ensure that UniControl can learn necessary and unique information from all tasks. For instance, if UniControl takes the segmentation map as the visual condition, the model might ignore the 3D information. As a result, the feature map learned may not be suitable for the task that takes the depth map images as visual condition. The second property would allow the model to learn the shared knowledge across tasks, as well as the differences among them.

We introduce two novel designed modules, MOE-style adapter and task-aware HyperNet , that allows UniControl enjoys the above two properties. An overview of the model design for UniControl is in Fig. 2. We describe the detailed designs of these modules below.

MOE-Style Adapter. Inspired by the design of Mixture-of-Experts (MOEs) [52], we devise a group of convolution modules to serve as the adapter for UniControl to capture features of various low-level visual conditions. Precisely, the designed adapter module can be expressed as

where 1 ( · ) is the indicator function, I k c is the conditioned image from task k , and F ( i ) Cov1 , F ( i ) Cov2 are the convolution layers of the i -th module of the adapter. We remove the weights of the original MOEs since our designed adapter is required to differentiate various visual conditions. Meanwhile, naive MOE modules can not explicitly distinguish different visual conditions when the weights are learnable. Moreover, such task-specific MOE adapters facilitate the zero-shot tasks with explicit retrieval of the adapters of highly related pre-training tasks. Besides, the number of parameters for each convolution module is approximately 70K, which is computationally efficient.

background

Conv

Conv

Conv

*0.4

Conv

*0.6

Estimated

tasks weights

Figure 3: Illustration of MOE's behaviors under zero-shot scenarios. The left part shows the capacity of the MOE to generalize to hybrid task conditions, achieved through the integration of outputs from two pertinent convolution layers. The right part illustrates the ability of the MOE-style adapter to generalize to unseen tasks, facilitated by the aggregation of pre-trained tasks using estimated weights.

Task-Aware HyperNet. The task-aware HyperNet modulates the zero-convolution modules of ControlNet [21] with the task instruction condition c task . As shown in Figure 2, our hyperNet first projects the task instruction c task into task embedding with the help of CLIPText encoder. Then similar in spirit of style modulation in StyleGAN2 [53], we inject the task embedding into the trainable copy of ControlNet, by multiplying the task embedding to each zero-conv layer. In specific, the length of the embedding is the same as the number of input channels of the zero-conv layer, and each element scalar in the embedding is multiplied to the convolution kernel per input channel. We also show that our newly designed task-aware HyperNet can also efficiently learn from training instances and task supervision following a similar analysis as in ControlNet [21].

3.3 Task Generalization Ability

With the comprehensive pretraining on the MultiGen-20M dataset, UniControl exhibits zero-shot capabilities on tasks that were not encountered during its training, suggesting that Unicontrol possesses the ability to transcend in-domain distributions for broader generalization. We demonstrate the zeroshot ability of UniControl in the following two scenarios:

Hybrid Tasks Generalization. As shown in the left side of Fig. 3, We consider two different visual conditions as the input of UniControl, a hybrid combination of segmentation maps and human skeletons, and augment specific keywords "background" and "foreground" into the text prompts. Besides, we rewrite the hybrid task instruction as a blend of instructions of the combined two tasks such as "segmentation map and human skeleton to image".

Zero-Shot New Tasks Generalization. As shown in the right side of Fig. 3, UniControl needs to generate controllable images on a newly unseen visual condition. To achieve this, estimating the task weights based on the relationship between unseen and seen pre-trained tasks is essential. The task weights can be estimated by either manual assignment or calculating the similarity score of task instructions in the embedding space. The example result in Fig. 5 (d) is generated by our manually assigned MOE weights as 'depth: 0.6, seg: 0.3, canny: 0.1' for colorization. The MOE-style adapter can be linearly assembled with the estimated task weights to extract shallow features from the newly unseen visual condition.

4 Experiments

We empirically evaluate the effectiveness and robustness of UniControl. We conduct a series of comprehensive experiments across various conditions and tasks, utilizing diverse datasets to challenge the model's adaptability and versatility. Experimental setup, methodologies, and results analysis are provided in the subsequent sections.

4.1 Experiment Setup

Implementation. The UniControl is illustrated as Fig. 2 with Stable Diffusion, ControlNet, MOE Adapter, and Task-aware HyperNet consisting ∼ 1.5B parameters. MOE Adapter consists of parallel convolutional modules, each of which corresponds to one task. The task-aware HyperNet inputs the CLIP text embedding [22] of task instructions and outputs the task embeddings to modulate the weights of zero-conv kernels. We implement our model upon the ControlNet . We take the AdamW [54] as the optimizer based on PyTorch Lightning [55]. The learning rate is assigned as 1 × 10 -5 . Our full-version UniControl model is trained on 16 Nvidia-A100 GPUs with the batch size of 4, requiring ∼ 5, 000 GPU hours. We have also applied Safety-Checker as safeguards of results.

Human Skeleton to Image

HED to Image

Object Bbox to Image

Normal Surface to Image

Image Outpainting

Segmentation to Image

Condition

Single-task ControlNet

Ours-Multi

Condition

Single-task ControlNet

Ours-Multi

A woman sitting on a bench near a statue, checking her phone.

Hotel room scene, efficiency, queen bed, sofa, table/desk, TV, in brown tone room.

A guy is having fun skiing down the slope of the hill.

Sunrise engagement

A brown bear walking in its zoo enclosure

A sink in a peninsula in a kitchen.

Art Decoration For Bedroom: Yellow Bedroom Walls Of Yellow Bedroom Ideas

Data Collection. Since the training set of ControlNet is currently unavailable, we initiate our own data collection process from scratch and name it as MultiGen-20M . We use a subset of LaionAesthetics-V2 [56] with aesthetics ratings over six, excluding low-resolution images smaller than 512. This yields approximately 2.8 million image-text pairs. Subsequently, we process this dataset for nine distinct tasks across five categories (edges, regions, skeletons, geometric maps, real images):

Canny (2.8M) : Utilize the Canny edge detector [57] with randomized thresholds.

HED(2.8M) : Deploy the Holistically-nested edge detection [58] for robust boundary determination.

Depth (2.8M) : Employ the Midas [59] for monocular depth estimation.

Normal (2.8M) : Use the depth estimation results from the depth task to estimate scene or object surface normals.

Segmentation (2.8M) : Implement the Uniformer [60] model, pre-trained on the ADE20K [61] dataset, to generate segmentation maps across 150 classes.

Object Bounding Box (874K) : Utilize YOLO V4 [62] pre-trained on the COCO [63] dataset for bounding box labelling across 80 object classes.

Human Skeleton (1.3M) : Employ the pre-trained Openpose [64] model to generate human skeleton labels from source images.

Image Outpainting (2.8M) : Create boundary masks for source images with random masking percentages from 20% to 80%.

Further processings are carried out on HED maps using Gaussian filtering and binary thresholding to simulate user sketching. Overall, we amass over 20 million image-prompt-condition triplets. Task instructions were naturally derived from the respective conditions, with each task corresponding to a specific instruction, such as "canny edge to image" for the canny task. We maintain a one-toone correspondence between tasks and instructions without introducing variance to ensure stability during training. We have additionally collected a testing dataset for evaluation with 100-300 imagecondition-prompt triplets for each task. The source data is collected from Laion and COCO. We will open-source our training and testing data to contribute to the community.

Benchmark Models. The most straightforward comparison for UniControl comes from task-specific ControlNet models. Six tasks overlap with those presented in ControlNet, so their official models are chosen as baselines for these tasks. For fair comparison, we re-implement the ControlNet model (single task) using our collected data. Our unified multi-task UniControl is compared against these

Output

Input

Output

(d) Zero-Shot Task: Image Colorization Example Result

Input

Output

'Once Upon A Time: S03E19, 720p, 'kiss them goodbye''

'Wedding hair and makeup with warm tones and red lips'

Input

Output

'Background: Photograph, Toronto Wet, Foreground: Some girls in

colorful shirts standing by some pastries'

'Background: Old City of Tzfat, Foreground: A set of twins who are taking a bath'

Condition 1: Depth

Condition 2: Human

Skeleton

Ours Result

Condition 1:

Segmentation

Condition 2: Human

Skeleton

Ours Result

(a) Zero Shot Combination: Depth and Human Skeleton to Image

(b) Zero Shot Combination: Segmentation and Human Skeleton to Image

(c) Zero-Shot Task: Image Deblurring Example Result

(e) Zero-Shot Task: Image Inpainting Example Result

task-aware models for each task. We apply default sampler as DDIM [32] with guidance weight 9 and steps 50. All single-task models used for comparison are trained by 100K iterations and our multi-task model is trained around 900K with similar iterations for each task to ensure fairness. The efficiency and compact design of our proposed model are evident in its construction. The total size of UniControl is around 1.5B # params and a single task ControlNet+SDM takes 1.4B. In order to achieve the same nine-task functionality, a single-task strategy would require the ensemble of a SDM with nine task-specific ControlNet models, amounting to approximately 4.3B # params in total.

4.2 Visual Comparison

We visually compare different tasks (Canny, HED, Depth, Normal, Segmentation, Openpose, Bounding Box, and Outpainting) in Fig. 4. Our method consistently outperforms the baseline ControlNet model. This superiority is in terms of both visual quality and alignment with conditions or prompts.

For the Canny task, the results generated by our model exhibit a higher degree of detail preservation and visual consistency. The outputs of UniControl maintain a faithful reproduction of the edge information ( i.e., round table) compared to ControlNet. In the HED task, our model effectively captures the robust boundaries, leading to visually appealing images with clear and sharp edge transitions, whereas ControlNet results appear to be non-factual. Moreover, our model demonstrate a more subtle understanding of 3D geometrical guidance of depth maps and surface normals than ControlNet. The depth map conditions produce visibly more accurate outputs. In the Normal task, our model faithfully reproduces the normal surface information ( i.e., ski pole), leading to more realistic

and visually superior outputs. During the Segmentation, Openpose, and Object Bounding Box tasks, the produced images generated by our model are better aligned with the given conditions than that by ControlNet, ensuring a higher fidelity to the input prompts. For example, the re-implemented ControlNet-BBox misunderstands 'a woman near a statue', whereas our outputs exhibit a high degree of accuracy and detail. In the Outpainting task, our model demonstrates its superiority by generating reasonable images with smooth transitions and natural-looking textures. It outperforms the ControlNet model, which produces less coherent results - 'a bear missing one leg'. This visual comparison underscores the strength and versatility of our approach across a diverse set of tasks.

4.3 Quantitative Evaluation

User Study. We compare the performance of our method with both the released ControlNet model and the re-implemented single-task ControlNet on our training set. As shown in Fig. 6, our approach consistently outperforms the alternatives in all cases. In the HED-to-image generation task, our method significantly surpasses ControlNet. This superiority is even more pronounced in the depth and normal surface to image generation tasks, where users overwhelmingly favor our method, demonstrating its ability to handle complex geometric interpretations. When compared to the re-implemented single-task model, Fig. 7 reveals that our approach maintains a smaller advantage, yet it still demonstrates its benefits by effectively discerning image regions to guide content generation. Even in the challenging outpainting task, our model outperforms the baseline, highlighting its robustness and capacity to generalize.

Image Perceptual Metric. We evaluate the distance between our output and the ground truth image. As we aim to obtain

a structural similar image to the ground truth image, we adopt the perceptual metric in [65], where a lower value indicates more similar images. As shown in Tab. 2, UniControl outperforms ControlNet on five tasks, and obtains the same image distance to ControlNet on Segmentation.

Fréchet Inception Distance (FID). We've further conducted quantitative analysis with FID [66] to include more classic single-task-controlled methods such as GLIGEN [67] and T2I-adapter [68]. With a collection of over 2,000 test samples sourced from Laion and COCO, we've assessed a wide range of tasks covering edges (Canny, HED), regions (Seg), skeletons (Pose), and geometric maps (Depth, Normal). The Tab. 3 demonstrates that our UniControl consistently surpasses the baseline methods across the majority of tasks. Notably, UniControl achieves this while maintaining a more compact and efficient architecture than its counterparts.

Ablation Study. We've conducted an ablation study, specifically focusing on the MoE-Style Adapter and TaskHyperNet in Tab. 4 with FID scores reported as the previous part. It is noticeable that the full-version UniControl (MoE-Style Adapter + TaskHyperNet) significantly outperforms the ablations which demonstrates the superiority of proposed MoE-Style Adapter and TaskHyperNet.

4.4 Zero-shot Generalization

We further showcase the surprising capabilities of our method to undertake the zero-shot challenge of hybrid conditions combination and unseen tasks generalization.

Hybrid Tasks Combination. This involves generating results from two distinct conditions simultaneously. Our model's zero-shot ability is tested with combinations such as depth and human skeleton or segmentation map and human skeleton. The results are shown in Fig. 5 (a)-(b). When the background is conditioned on a depth map, the model effectively portrays the intricate 3D structure of the scene, while maintaining the skeletal structure of the human subject. Similarly, when the model is presented with a combination of a segmentation map and human skeleton, the output skillfully retains the structural details of the subject, while adhering to the segmentation boundaries. These examples illustrate our model's adaptability and robustness, highlighting its ability to handle complex hybrid tasks without any prior explicit training.

Unseen Tasks Generalization. To evaluate the zero-shot ability to generalize to unseen tasks such as gray image colorization, image deblurring, and image inpainting, we conduct the case analysis in Fig. 5 (c)-(e). The model skillfully handles the unseen tasks, producing compelling results. This capability is deeply rooted in the shared attributes and implicit correlations among pre-training and new tasks, allowing our model to adapt seamlessly. For instance, the colorization task leverages the model's understanding of image structures from the segmentation task and depth estimation task, while deblurring and inpainting tasks benefit from the model's familiarity with edge detection and outpainting ones.

5 Conclusion and Discussion

We introduce UniControl , a novel unified model for incorporating a wide range of conditions into the generation process of diffusion models. UniControl has been designed to be adaptable to various tasks through the employment of two key components: a Mixture-of-Experts (MOE) style adapter and a task-aware HyperNet. The experimental results have showcased the model's robust performance and adaptability across different tasks and conditions, demonstrating its potential for handling complex text-to-image generation tasks.

Limitation and Broader Impact. While UniControl demonstrates impressive performance, it still inherits the limitation of diffusion-based image generation models. Specifically, it is limited by our training data, which is obtained from a subset of the Laion-Aesthetics datasets. We observe that there is a data bias in this dataset. Although we have performed keywords and image based data filtering methods, we are aware that the model may generate biased or low-fidelity output. Our model is also limited when high-quality human output is desired. UniControl could be improved if better open-source datasets are available to block the creation of biased, toxic, sexualized, or other harmful content. We hope our work can motivate researchers to develop visual generative foundation models.

References

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 21(140):1-67, 2020.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416 , 2022.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35:27730-27744, 2022.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

OpenAI. Gpt-4 technical report, 2023.

Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping languageimage pre-training with frozen image encoders and large language models. arXiv preprint arXiv:2301.12597 , 2023.

Haoxuan You, Mandy Guo, Zhecan Wang, Kai-Wei Chang, Jason Baldridge, and Jiahui Yu. Cobit: A contrastive bi-directional image-text generation model. arXiv preprint arXiv:2303.13455 , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In CVPR , pages 10684-10695, 2022.

Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. arXiv preprint arXiv:2304.02643 , 2023.

Bo Li, Dongseong Hwang, Zhouyuan Huo, Junwen Bai, Guru Prakash, Tara N Sainath, Khe Chai Sim, Yu Zhang, Wei Han, Trevor Strohman, et al. Efficient domain adaptation for speech foundation models. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 1-5. IEEE, 2023.

Rongjie Huang, Mingze Li, Dongchao Yang, Jiatong Shi, Xuankai Chang, Zhenhui Ye, Yuning Wu, Zhiqing Hong, Jiawei Huang, Jinglin Liu, et al. Audiogpt: Understanding and generating speech, music, sound, and talking head. arXiv preprint arXiv:2304.12995 , 2023.

Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691 , 2022.

Yilun Dai, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.00111 , 2023.

Sai Vemprala, Rogerio Bonatti, Arthur Bucker, and Ashish Kapoor. Chatgpt for robotics: Design principles and model abilities. Technical Report MSR-TR-2023-8, Microsoft, February 2023. URL https://www.microsoft.com/en-us/research/publication/ chatgpt-for-robotics-design-principles-and-model-abilities/ .

Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155 , 2022.

OpenAI. Chatgpt. https://openai.com/blog/chatgpt/ , 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. NeurIPS , 34:8780-8794, 2021.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 , 2022.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543 , 2023.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. arXiv preprint arXiv:2103.00020 , 2021.

Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3712-3722, 2018.

Golnaz Ghiasi, Barret Zoph, Ekin D Cubuk, Quoc V Le, and Tsung-Yi Lin. Multi-task selftraining for learning general representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 8856-8865, 2021.

Shikun Liu, Edward Johns, and Andrew J Davison. End-to-end multi-task learning with attention. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 1871-1880, 2019.

David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks. In International Conference on Learning Representations , 2017.

Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.00695 , 2019.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML , pages 2256-2265. PMLR, 2015.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurIPS , 33:6840-6851, 2020.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv:2010.02502 , October 2020. URL https://arxiv.org/abs/2010.02502 .

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI , pages 234-241. Springer, 2015.

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res. , 23(47):1-33, 2022.

Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:2112.10741 , 2021.

Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. In European Conference on Computer Vision , pages 423-439. Springer, 2022.

Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In CVPR , 2017.

Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACMSIGGRAPH 2022 Conference Proceedings , pages 1-10, 2022.

Tengfei Wang, Ting Zhang, Bo Zhang, Hao Ouyang, Dong Chen, Qifeng Chen, and Fang Wen. Pretraining is all you need for image-to-image translation. arXiv preprint arXiv:2205.12952 , 2022.

Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations , 2021.

Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt-to-prompt image editing with cross attention control. arXiv preprint arXiv:2208.01626 , 2022.

Bram Wallace, Akash Gokul, and Nikhil Naik. Edict: Exact diffusion inversion via coupled transformations. arXiv preprint arXiv:2211.12446 , 2022.

Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual diffusion implicit bridges for image-to-image translation. In The Eleventh International Conference on Learning Representations , 2022.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 2022.

Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.14699 , 2022.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. I 2 sb: Image-to-image schrödinger bridge. arXiv preprint arXiv:2302.05872 , 2023.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2022.

Zhendong Wang, Yifan Jiang, Yadong Lu, Yelong Shen, Pengcheng He, Weizhu Chen, Zhangyang Wang, and Mingyuan Zhou. In-context learning unlocked for diffusion models. arXiv preprint arXiv:2305.01115 , 2023. URL https://arxiv.org/abs/2305.01115 .

Amir Bar, Yossi Gandelsman, Trevor Darrell, Amir Globerson, and Alexei Efros. Visual prompting via image inpainting. Advances in Neural Information Processing Systems , 35: 25005-25017, 2022.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 , 2017.

Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 8110-8119, 2020.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurIPS . NeurIPS, 2019.

Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2111.02114 , 2021.

John Canny. A computational approach to edge detection. IEEE Transactions on pattern analysis and machine intelligence , pages 679-698, 1986.

Saining Xie and Zhuowen Tu. Holistically-nested edge detection. In Proceedings of the IEEE international conference on computer vision , pages 1395-1403, 2015.

René Ranftl, Katrin Lasinger, David Hafner, Konrad Schindler, and Vladlen Koltun. Towards robust monocular depth estimation: Mixing datasets for zero-shot cross-dataset transfer. IEEE transactions on pattern analysis and machine intelligence , 44(3):1623-1637, 2020.

Kunchang Li, Yali Wang, Junhao Zhang, Peng Gao, Guanglu Song, Yu Liu, Hongsheng Li, and Yu Qiao. Uniformer: Unifying convolution and self-attention for visual recognition. arXiv preprint arXiv:2201.09450 , 2022.

Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision , 127:302-321, 2019.

Alexey Bochkovskiy, Chien-Yao Wang, and Hong-Yuan Mark Liao. Yolov4: Optimal speed and accuracy of object detection. arXiv preprint arXiv:2004.10934 , 2020.

Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13 , pages 740-755. Springer, 2014.

Zhe Cao, Tomas Simon, Shih-En Wei, and Yaser Sheikh. Realtime multi-person 2d pose estimation using part affinity fields. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 7291-7299, 2017.

Richard Zhang, Phillip Isola, Alexei Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2018.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. NeurIPS , 30, 2017.

Yuheng Li, Haotian Liu, Qingyang Wu, Fangzhou Mu, Jianwei Yang, Jianfeng Gao, Chunyuan Li, and Yong Jae Lee. Gligen: Open-set grounded text-to-image generation. arXiv:2301.07093 , 2023.

Chong Mou, Xintao Wang, Liangbin Xie, Jian Zhang, Zhongang Qi, Ying Shan, and Xiaohu Qie. T2i-adapter: Learning adapters to dig out more controllable ability for text-to-image diffusion models. arXiv preprint arXiv:2302.08453 , 2023.

A Details of Implementation

A.1 MOE-Style Adapter

The MOE adapter is implemented as a set of parallel ConvNets composed of three consecutive convolution and non-linear activation layers. The entire model is comprised of nine individual MOE adapters, each of which consumes 70K parameters. Task keys are designated to each adapter, ensuring that they align with the corresponding visual conditions. Once the MOE adapter processes the input, the remaining model parameters become shared across all tasks. This architecture facilitates task adaptability while promoting parameter efficiency.

A.2 Task-aware HyperNet

The task-aware hypernet is applied to modulate the parameters of zero-conv layers in the ControlNet. Since the ControlNet can be considered as the hypernet of Stable Diffusion (fixed copy). Our idea can be concluded as the control over control or meta-control to let the task-aware hypernet learn the universe representation that is generalizable across different tasks. To implement it, we firstly map the task keys to instruction with a mapping function as: {"hed": "hed edge to image", "canny": "canny edge to image", "seg": "segmentation map to image",

"depth": "depth map to image", "normal": "normal surface map to image",

"pose": "human pose skeleton to image", "hedsketch": "sketch to image",

A.3 Data Collection

We have collected a large amount of training set (MultiGen-20M) including over 20M conditionimage-prompt triplets across nine different tasks. We firstly download 3/4 of Laion-Aesthetics-V2 with score over six and filter out low-resolution (<512) images. As a result, 2.8M images are selected as source images. Then we apply the visual condition extractors as described in the main paper to collect Canny , HED , Sketch , Depth , Normal Surface , Seg Map , Object Bounding Box , Human Skeleton and Outpainting .

B Numerical Analysis of Task-Aware Modulated ControlNet

We show that our proposed task-aware modulated ControlNet preserves the properties of the original ControlNet structure. Specifically, we show 1) The new task-aware modulated ControlNet preserves the zero-initialization property of ControlNet; 2) The parameters of the task-aware modulated Controlnet can be updated once we start to train the model.

Denote the input feature map by x , the frozen SD Block in Fig. 2 by F SD , the extra condition by c , two zero convolution operators by Z 1 θ 1 ( · ) and Z 2 θ 2 ( · ) , the trainable copy of SD Block by G SD θ s ( · ) , the task instruction by c task , and the task-aware hyperNet by H θ H ( · ) . Then the output of the new task-aware modulated Controlnet can be expressed as

Property of Zero Initialization. Similar to ControlNet [21], the weights and biases of the convolution layers are initialized as zeros. As a result, we have Z 1 θ 1 ( · ) ≡ 0 and y c = F SD ( x ) , regardless of the initialization of H θ H ( · ) .

Appendix

Gradient Analysis. We analyze the gradient of the modulated part

where I is the input of the zero convolution layer.

When we start to train the network, the first part of the RHS of (2) follows similar analysis of ControlNet [21] since H θ H ( c task ) is constant when we analyze the gradient ∇ θ Z 1 θ 1 ( I ) . Since the parameters of H θ H ( c task ) are not initialized to zero, it is known that H θ H ( c task ) = 0 . So the gradient dynamic follows the analysis of ControlNet. Therefore, we conclude that Z 1 θ 1 ( I ) = 0 after the first gradient update, and that the network can start to learn and update the following standard dynamics of stochastic gradient descent.

̸

̸

As for the second part of the RHS of (2), Z 1 θ 1 ( I ) ≡ 0 before the first gradient update, so the gradient is zero for θ H . However, after the first gradient update of θ 1 , we know Z 1 θ 1 ( I ) = 0 , and θ H can be updated with non-zero gradients.

To conclude, the new task-aware Modulated ControlNet can still be efficiently updated and learned even if the convolution layers are initialized to zero.

C Zero-shot-task Results and Analysis

We show more zero-shot-task results in this section, where the tasks have not been trained on. In Fig. 13, we show zero-shot deblurring results guided by the keywords. Our deblurred images can successfully recover the fine-grained details of the images without training on such data. We note that some details are still missing, e.g., the details in the painting in the first row are still not clear enough. In Fig. 14, we illustrate two zero-shot image colorization results. We believe that most parts of the generated images are acceptable, though the clothes of the second woman do not look the same to the input blurred image. In Fig. 15, we observe impressive zero-shot inpainting results. In the first row, the duck that is inputted in the text has been successfully generated in the inpainted image. The second row obtains acceptable results as well, though the faces do not look perfect. The overall zero-shot quality of UniControl is remarkable.

While inpainting and outpainting might appear related, they are fundamentally distinct. Inpainting heavily leverages the contextual information from unmasked regions, necessitating a precise match. Conversely, outpainting has more freedom, with the generative model prioritizing prompts to envision new content. As shown in Fig. 8, directly using outpainting model for inpainting tasks can be challenging since the model tends to leave a sharp change over the mask boundaries. Our pretrained UniControl, thanks to intensive training across multiple tasks, has learned edge and region-to-image mappings, which assists in preserving contextual information.

Our model also demonstrates a promising capacity to generalize under scribble conditions, showing parallels to the ControlNet's ability, even though UniControl hasn't been directly trained using scribble data. Fig. 9 provides results illustrating the scribble-to-image generation.

D Details of User Study

In the evaluation steps, we use Amazon Mechanical Turk (Mturk) 2 to perform user study. Specifically, we ask three Mturk master workers to select the best output result for each input condition. As shown in Fig. 10, we provide instructions on guidelines to select the best generated image. The annotators are provided the condition map and the text that describes the image, and are required to select the better output between the two generated images. Considering that images can both in good or bad qualities, we provide the tie option as well. We use the majority vote to determine the result of each image, which means that an image is considered as a better image if two or more annotators vote for it. We use 294 images for the tasks of Canny, HED, Surface Normal, Depth, Segmentation, User Sketch, and Outpainting. We adopt 100 images for the task of Human Skeleton and 187 images for the task of Bounding Box. In summary, we totally obtain 7,035 voting results for all nine tasks. 2/3 of source images in testing set are collected from MSCOCO with the remaining 1/3 from Laion. And it includes a very diverse range of topics including indoor scene, outdoor scene, oil painting, portrait, pencil sketch, animation, cartoon, etc.

̸

Input (User Scribble)

Input (User Scribble)

Input (User Scribble)

UniControl

Default

Automatic Prompt

User Prompt

Default

Automatic Prompt

User Prompt

'

a turtle in river

'

'

a masterpiece of cartoon-style turtle illustration

'

'

a cow with horns standing in a field

'

'

a robot ox on moon, UE5 rendering, ray tracing

'

'

a digital painting of a hot air balloon

'

'

magic hot air balloon over a lit magic city at night

'

'

a turtle in river

'

'

a masterpiece of cartoon-style turtle illustration

'

'

a cow with horns standing in a field

'

'

a robot ox on moon, UE5 rendering, ray tracing

'

'

a digital painting of a hot air balloon

'

'

magic hot air balloon over a lit magic city at night

'

Default

Automatic Prompt

User Prompt

'

a turtle in river

'

'

a masterpiece of cartoon-style turtle illustration

'

'

a cow with horns standing in a field

'

'

a robot ox on moon, UE5 rendering, ray tracing

'

Input (Scribble)

ControlNet-Scribble Results

UniControl Results (zero-shot)

'a turtle in river'

'a digital painting of a hot air balloon'

'

a door on a wall

'

'

magical door, Hearthstone

'

'

a door on a wall

'

'

magical door, Hearthstone

'

'

a digital painting of a hot air balloon

'

'

magic hot air balloon over a lit magic city at night

'

'a door on a wall'

'

'

'

a door on a wall

'

an elephant with background in the field an elephant with background in the field

'

'

'

'

'

magical door, Hearthstone

'

Egyptian elephant sculpture

Egyptian elephant sculpture

'

'

Instructions

Shortcuts

Allimages are generated by Al based on the same text input and depth condition, but using different Al methods. A depth edge map is the depth of objects in the real world, which can be obtained by Lidar. Please select the image with the BEST image qualty that aligns to the text an...

?

DepthImage

Generated Image1

Generated Image 2

Generated Image 2 is better 2

Imagetitle:camponamountain top:Birthday1

2

Presents,Adventure,Outdoor,Mountain

Equally good or bad

Camps,GreatView,Places,Hiking,Mornings

Allimaaes_are_edited.hy Alhased.on.the_eame inetruction_hutueina different.methodsWe_nrovide_a.denth_edoe_man.to.auide_the_imaae

Submit

Ours-Multi

Ours-Single

Tie

0

20

40

60

80

76.2%

17%

6.8%

Sketch-to-Image

E Failure Cases

We illustrate some failure cases in Fig. 12. In the first row, although our generated image successfully aligns the Bounding Box condition, the generated human has a distorted body. In the second row, our generated image looks similar to the ground truth; however, the human faces are blurred. We think that the reason is that UniControl inherits the data and model bias of Stable Diffusion, where the generated human commonly have issues. In the third row, the generated image does not look realistic. We believe that the training data can be improved both quantitatively and qualitatively.

F Additional Results

We illustrate more visualized results in this section on tasks Canny (Fig. 16), HED (Fig. 17), Depth (Fig. 18), Surface Normal (Fig. 19), Human Skeleton (Fig. 20), Bounding Box (Fig. 21), Segmentation (Fig. 22) and Outpainting (Fig. 23). These results further demonstrate the effectiveness of our proposed method. Moreover, due to the space limitation in the main paper, we report results of the last task, User Sketch. Given a sketched image, UniControl is able to achieve promising realistic images. The visualized results are in Fig. 24. The user study result can be found in Fig. 11, where it is observed that UniControl obtains significantly more votes than the single task model.

e

Blurred Image

国

Our Result

OO0

ADGEE

Gray Image

Our Result

Cropped Image

Our Result

10

10

'Contemporary Bedroom Designs 2015 modern bedroom designs intended design '

<Caption>: two arms typing on a laptop and one hand on a mouse

Our Method Output

<Caption>: Two people walking along a side walk next to a train on the tracks.

Our Method Output

<Caption>: A close up of glazed donuts that are plain or with chocolate.

(c) 'A close up of glazed donuts that are plain or with chocolate.'

Our Method Output

<Caption>: A group of elephants with water in front and trees behind.

(d) 'A group of elephants with water in front and trees behind.'

山

<Caption>: a person on skis makes her way through the snow

<Caption>: three zebras grazing in a grassy area near shrubs

<Caption>: The Taj Mahal mirrored by a water fountain's reflection. - Agra, Uttar Pradesh, India - Daily Travel Photos

<Caption>: A young girl who is brushing her teeth with a toothbrush.

(c) 'The Taj Mahal mirrored by a water fountain's reflection. - Agra, Uttar Pradesh, India - Daily Travel Photos'

(d) 'A young girl who is brushing her teeth with a toothbrush.'

<Caption>: The Ta Prohm Temple Located at Angkor in Cambodia by Kyle Hammons

<Caption>: A brown dog standing on a wooden bench near a lemon tree.

<Caption>: A Polar Bear walks toward water, while a large bird lands on the opposite bank.

<Caption>: A display of vintage animal toys on the floor.

AQSHAUW

DRAVIAUNS

<Caption>: A line of small teddy bears are in front of several DVD cases.

69

alao Ooco

<Caption>: A man in the kitchen standing with his dog.

<Caption>: A hot dog sitting on top of a bun in a wrapper.

Our Method Output

Our Method Output

(d) 'Several people waiting on the side of train tracks as a train with it's lights on comes down the track'

<Caption>: Photo of handsome man in black leather jacket

不

<Caption>: A woman is sitting near a prominent landmark.

<Caption>: A man that has ski's and is standing in the snow.

<Caption>: a woman is sitting in front of a desk

A

<Caption>: A woman is walking two dogs in the snow.

1

Our Method Output

(a) 'A woman is walking two dogs in the snow' <Caption>: Simone Righi frasi glasses linen suit menswear streetstyle icon fashion florence.jpg

<Caption>: The large room has a wooden table with chairs and a couch.

Our Method Output

<Caption>: two black bags placed standing on the ground

Our Method Output

(c) 'The large room has a wooden table with chairs and a couch.'

Our Method Output

(d) 'Two black bags placed standing on the ground'

<Caption>: A bench at the beach next to the sea

一

<Caption>: A well-lit and well-decorated living room shows a glimpse of a glass front door through the corridor.

<Caption>: Blue Hour Barley

Our Method Output

(d) 'Blue Hour Barley'

<Caption>: St, Michaels Mount, Cornwall

<Caption>: Melbourne by teekay 72

<Caption>: Lady in Black Kimono Paint by Diamonds

<Caption>: Beautiful kitchen grand scale living pinterest for Kitchen cabinets lowes with old world metal wall art

(d) 'Beautiful kitchen grand scale living pinterest for Kitchen cabinets lowes with old world metal wall art'

<Caption>: A Limited Edition, Fine Art photograph of a beautiful sunrise at Lake Jackson in Sebring, Florida. Available as a Fine Art print

<Caption>: Giraffes in Lake Manyara national park

(a) 'A Limited Edition, Fine Art photograph of a beautiful sunrise at Lake Jackson in Sebring, Florida. Available as a Fine Art print'

<Caption>: Wild Ones 1991 Limited Edition Print - Frank McCarthy

Cla

<Caption>: Superhero watching over city. No transparency used. Basic (linear) gradients. A4 proportions.

(d) 'Superhero watching over city. No transparency used. Basic (linear) gradients. A4 proportions.'

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a unified diffusion framework, UnitControl for a more fine-grained controllable generation based on input visual conditions, utilizing two modules mixture-of-experts adaptor that extract features from different visual conditions, and a task-aware HyperNet that extracts language-based task embedding to control visual generation.\n\nSTRENGTHS:\n1. The paper conducts a relatively comprehensive evaluation of the proposed method, including human evaluations.\n2. The framework can be useful as it unified different modalities for visual controllable generation using diffusion models.\n3. The authors promise to open-source 20M triple training datasets that can benefit the whole community.\n\nWEAKNESSES:\n1. **The writing is not so well written and missing details that can hinder understanding of readers**. For example, Figure 2 should also mention parameters of encoder and decoder since I initially thought the introduced parameters are only from MOE adaptors and Hypernet. The authors should also highlight the differences and **similarities** (e.g., encoder and decoder are copied) between their work and ControlNet. In Section 3.3, how is hybrid task generalization done? Since your model only conditions on one visual modality when computing the score in one forward pass, do you compute scores individually for each condition for generation? If so, what sampling technique did you use? One sampling technique I can think of is the conjunction operator from Composable Diffusion [1] or do you use some other existing techniques? For zero-shot task generalization, I also find it confusing and don't quite understand how it is done. Mentioning it in the supplementary is not sufficient since it is meant to provide details that don't affect your understanding overall.\n\n2. **The claims are quite shaky.** The authors claim that the model need to tackle the misalignment of low-level features from different tasks. However, based on Figure 4, I don't see too much of misalignment from baselines, i.e., ControlNet. Besides the visual quality, I don't see much of advantages over ControlNet. Even sometimes single-task ControlNet is better in some aspects. For example, for the object bounding box to image, control net seems to understand there should be two separate benches instead of one long bench generated by the proposed model. In addition, since one unified model is trained for longer time, model could be more robust, directly improving overall visual quality.\n\n3. **The zero-shot generalization ability could be misleading.** Suppose in training, there doesn't exist any instructions that perform colorization. I don't expect the language itself can bridge the gap between visual signals and language to perform such unseen tasks. For example, in Figure 5, such zero-shot results doesn’t necessarily reflect the true reality that the model understands given instructions. One way to check this is to reverse the whole process - for example, can you do image decolorization using your prompt or blur the image instead. Since training images are mostly high resolution, and colored images, the model could automatically utilize learned prior to make images higher quality, which doesn’t necessarily mean they are doing what they are told to do. If so, then the model just does general reconstruction when prompted something unknown.\n\n4. **Missing quantitative evaluations.** As stated, one contribution is to show that the visual quality of such controllable generation can be improved, then it is needed to include quantitative comparisons for image fidelity using metrics such as FID and KID, though I don't think they are good metrics but it provides enough insights.\n\n5. **Lacking relevant baselines.** There are also other existing adaptor methods for such generation, for example, T2I-adaptor [2] and GLIGEN [3] seem to be highly relevant, and both of them are not used as a comparison in the paper.\n\n5. **Related work.** If one of the contributions is to combine different modalities for compositional generation, then it would be important to add existing works related to compositional generation.\n\n[1] Liu et al., Compositional Visual Generation with Composable Diffusion Models (ECCV 2022) \\\n[2] Mou et al., T2I-Adapter: Learning Adapters to Dig out More Controllable Ability for Text-to-Image Diffusion Models. \\\n[3] Li et al., GLIGEN: Open-Set Grounded Text-to-Image Generation (CVPR 2023)\n\nQUESTIONS:\n1. How are both hybrid tasks generalization and zero-shot new task generalization done? For example, a detailed procedure (i.e., sampling technique) will be useful for better understanding.\n2. There are 6 tasks overlapped with ControlNets, each of which is trained for 100K iterations. However, there are 8 tasks in your multi-task model, but it is trained for 900K, which in fact is not a fair comparison. If I misunderstood, feel free to point that out.\n3. The method mainly uses modalities that are spatial, so is it capable of using global context, e.g., color, texture, to guide image generation?\n\nLIMITATIONS:\nThe authors have included limitations and broader impact such as training data bias.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUniControl is a diffusion-based image generation model that can condition on natural language input as well as multiple types of visual inputs (e.g. edge map, depth map). The framework is built upon components of Stable Diffusion Models, an MOE adapter and a ControlNet modulated by a task-aware HyperNet. The model demonstrates abilities of conditioning on one or multiple visual inputs at a time, as well as visual inputs that it has not seen during training.\n\nSTRENGTHS:\n- UniControl extends ControlNet to work with multiple tasks and shows that the tasks help each other so as to improve performance on single task metrics as well.\n- the MOE adapter set up easily allows the model to condition on multiple visual inputs.\n\nWEAKNESSES:\n- The model has a certain complexity, as it involves multiple modules such as the MOE adapter, the ControlNet and a HyperNet.\n- It needs to maintain two sets of SDM parameters. (This is the same as ControlNet, though.)\n- MOE adapter means more parameters to be added with each new task added.\n- Task instruction is needed for both training and inference, and is handled by a separate module than the language prompt, even though both are text prompts. This can be a downside if the task instruction is unknown or not well defined.\n- It is not clear what $c_{task}$ is used for the task-aware hypernet during hybrid-task or zero-shot new task inference.\n\nAblation study would be helpful to show the importance/usefulness of the proposed components:\n- model performance without the MOE adapter.\n- model performance without conditioning on $c_{task}$, or merge $c_{task}$ into $c_{text}$ (in a thoughtful way). This will show whether the task-aware hypernet is needed.\n\nQUESTIONS:\n- HyperNet should be cited?\n\nLIMITATIONS:\nThe authors have discussed the limitations and broader impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper introduces UniControl, a new generative foundation model that consolidates a wide array of controllable condition-to-image (C2I) tasks within a singular framework. UniControl enables pixel-level-precise image generation, where visual conditions primarily influence the generated structures and language prompts guide the style and context. For this purpose, the authors augment pretrained text-to-image
diffusion models (ControlNet) and introduce a task-aware HyperNet to modulate the diffusion models, enabling the adaptation to different C2I tasks simultaneously. UniControl was trained on nine unique C2I task and demonstrated excellent zero-shot generation abilities with unseen visual conditions.

STRENGTHS:
1. The paper is clearly written and easy to follow.
2. The related work section covers the most relevant papers in the field.
3. The approach produces excellent image generation results.
4. The experimental evaluation is convincing.

WEAKNESSES:
1. It relies too much on existing image generation methods (ControlNet).
2. A more scientific contribution would have been expected.

QUESTIONS:
Here are my concerns:
1. The paper claim in the abstract that: "...UniControl often surpasses the performance of single-task-controlled methods of comparable model sizes". Maybe I missed something, but could the authors elaborate more on this statement? What are the other 'single-task-controlled methods' they refer to? I only found comparison with ControlNet.
2. Could you adapt your approach to work with other Image Generation software (besides Stable Diffusion)?

LIMITATIONS:
Limitations are addressed in the paper in a dedicated section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a method for controlling the output of a diffusion\nmodel with multiple modalities of reference images, e.g. edges,\nsegmentation, depth, etc. It can be seen as proposing a multi-task\nversion of ControlNet. Experiments in the paper show the multi-task\napproach outperforms the single-task, and also allows zero-shot\napplicability to novel tasks such as colorization.\n\nSTRENGTHS:\n- The paper effectively demonstrates architecture modifications for multi-task controlling model\n  and shows that it is generally better than the single-task (Table\n  1).\n- The model generalization to tasks that were not in the training set such\n  as colorization, deblurring or inpainting is quite remarkable.\n- The paper introduces a dataset with 20M multi-modal condition\n  training pairs.\n\nWEAKNESSES:\nThe paper presents a well engineered solution to achieving a\nmulti-task version of ControlNet, and show generalization to some new\ntasks. Some decision justifications are not backed up by ablation\nexperiments and the experiments on task generalization are\ndemonstrated only with a few visual examples.\n\n- If I understood correctly, there \"Mixture of Experts\" component is\n  manually selecting the encoder for each modality. In that case\n  calling this module a MOE is justified, since it could be misleading.\n- Unless I missed it, the different components in the design are not\n  ablated. What is the contribution of the Hypernet task embedding?\n- How do you explain the poor performance of ControlNEt in\n  Normal-to-Image (Fig. 6)\n- Quantitative comparison for the generalization to new tasks is\n  limited to a few visual examples. Some of the examples are even\n  repeated three times (segmentation+skeleton); for that case, it may\n  be more convincing if the paper showed three different results.\n\nQUESTIONS:\nPlease refer to the weaknesses section.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents UniControl, that unifies multiple visual controlling condition into a single unified model. To achieve this, the authors introduce a task-aware HyperNet to modulate the diffusion models, enabling adaptation to different condition-to-image (C2I) tasks simultaneously. UniControl is trained on nine unique C2I tasks and demonstrates impressive controlled generation quality. It also demonstrates zero-shot generation abilities with unseen visual conditions, including condition combination and new conditioning. User studies show that UniControl often surpasses the performance of single-task-controlled methods of comparable model sizes.\n\nSTRENGTHS:\n- The paper addresses an interesting and important research topic, unifying the controllability of different models into a single model.\n- The authors compare the results with ControlNet and demonstrate better qualitative results.\n- The zero-shot task generalization and instruction combination aspects of the proposed method are intriguing and valuable.\n\nWEAKNESSES:\n- The paper lacks quantitative evaluations for the alignment of the generated content and conditional inputs. For segmentation and bounding box, pretrained detectors could be used for evaluating the alignment, following the configuration in [1]. Although ControlNet does not handle bounding box conditions, a baseline for object detection instruction-following would be [1].\n\n[1] Li, Y., Liu, H., Wu, Q., Mu, F., Yang, J., Gao, J., ... & Lee, Y. J. (2023). Gligen: Open-set grounded text-to-image generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 22511-22521).\n\n- Zero-shot task generalization seems underexplored, and there is a lack of detailed analysis.\n    - Inpainting results are not surprising as they share a strong alignment in terms of task instruction with the pretraining with outpainting.\n    - For deblurring and colorization, it would be better to illustrate how the model can perform zero-shot task generalization. Does the capability mainly come from the task instruction encoder?\n    - If so, it would be more intriguing to demonstrate results with more tasks that can really benefit from the zero-shot task generalization. It is easy to obtain data for colorization, inpainting. It would be valuable to demonstrate results on tasks that may be hard to collect training data.  For example, for scribble to image, ControlNet utilizes strong human-crafted data augmentation to synthesize the scribbles. Can UniControl generalize well to this case?\n\nQUESTIONS:\n- Can the authors please clarify if the difference between Ours-Single and ControlNet is only the training data / training schedule or there are other differences? From the results in Fig. 6 and Fig. 7, it seems that Ours is much better than ControlNet, while Ours is only slightly better than Ours-Single.\n\n- Where does the performance can mainly come from ControlNet to Ours-Single? Is it because of the better dataset and/or longer training? Note that this is not a criticism, so I put it here in the questions. But I think a clarification or analysis in this regard can be helpful for the readers and the research community.\n\n- Cost comparison with ControlNet. In L38-39, the authors mention \"Retraining a separate model is necessary to handle a different modality of visual conditions, incurring non-trivial time and spatial complexity costs.\" However, UniControl requires 5000 GPU hours on A100, while ControlNet training for a single model is usually just 100-300 hours. This somehow invalidates this point.\n\nLIMITATIONS:\nThe authors discuss the limitations in the main papar.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}