SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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*, and Ran Xu†
†Salesforce AI Research, †Northeastern University, ‡Stanford Univeristy, qin.ca@northeastern.edu, ermon@cs.stanford.edu, yunfu@ece.neu.edu, {shu.zhang, ning.yu, yihaof, x.yang, yingbo.zhou, huan.wang, jniebles, cxiong, ssavarese, ran.xu}@salesforce.com
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}$
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
$^{1}$ Code: https://github.com/salesforce/UniControl
UniControl Diffusion Model
Pre-training Tasks
Canny Edge
to Image
HED Edge to
Skeleton to
Outpainting
Zero-shot Combination
Segmentation +
Skeleton to Image
Depth Map +
Zero-shot Tasks
Gray Image Colorization
Image Deblur
Zero-shot Task: Image Deblurring
Zero-shot Task: Image Colorization
"A cave on a desert planet"
Zero-shot Task: Seg-and-Skeleton-to-Image
Zero-shot Task: Depth-and-Skeleton-to-Image
Figure 1: UniControl is trained with multiple tasks with a unified model, and it further demonstrates promising capability in zero-shot tasks generalization with visual example results shown above.
Input Condition
"Chocolate mousee cake"
"Casual Short Hairstyles For Women"
"One tray of rice and a tray of fruits and veggies."
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.
Table 1: Architecture and Model Size (#Params): UniControl vs. Multi-ControlNet
Stable Diffusion	ControlNet	MoE-Adapter	TaskHyperNet	Total
UniControl	1065.7M	361M	0.06M	12.7M	1.44B
Multi-ControlNet	1065.7M	361M × 9	-	-	4.32B
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 image-text-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.
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 text-guided 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 diffusion-based 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.
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:  $\mathcal{D} := \{\mathcal{D}_1 \cup \dots \cup \mathcal{D}_K\}$ , and for each task training set  $\mathcal{D}_k$ , denote the training pairs by  $([c_{\mathrm{text}}, c_{\mathrm{task}}], \mathcal{I}_c, \pmb{x})$ , with  $c_{\mathrm{task}}$  being the task instruction that indicates the task type,  $c_{\mathrm{text}}$  being the language prompt describing the target image,  $\mathcal{I}_c$  being the visual conditions, and  $\pmb{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:
where the task is to translate the canny edge to real images following language prompt. With the induced training pairs  $(\pmb{x},[c_{\mathrm{task}},c_{\mathrm{text}}],\mathcal{I}_c)$ , we define the training loss for task  $k$  following LDM [9]:  $\ell^k (\theta)\coloneqq \mathbb{E}_{z,\varepsilon ,t,c_{\mathrm{task}},c_{\mathrm{text}},\mathcal{I}_c}\left[\| \varepsilon -\varepsilon_\theta (z_t,t,c_{\mathrm{task}},c_{\mathrm{text}},\mathcal{I}_c)\| _2^2\right]$  with  $([c_{\mathrm{task}},c_{\mathrm{text}}],\mathcal{I}_c,\pmb {x})\sim \mathcal{D}_k$  where  $t$  represents the time step,  $z_{t}$  is the noise-corrupted latent tensor at time step  $t$ ,  $z_0 = E(\pmb {x})$  and  $\theta$  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  $\mathcal{D}_k$  and optimize  $\theta$  with the calculated loss  $\ell^k (\theta)$ .
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
Figure 2: This figure shows our proposed UniControl method. To accommodate diverse tasks, we've designed a Mixture of Experts (MOE) Adapter, containing roughly  $\sim 70\mathrm{K}$  #params for each task, and a Task-aware HyperNet ( $\sim 12\mathrm{M}$  #params) to modulate  $N$  (i.e., 7) zero-conv layers. This structure allows for multi-task functionality within a singular model, significantly reducing the model size compared to an equivalent stack of single-task models, each with around 1.4B #params.
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
$$
\mathcal {F} _ {\mathrm {A d a p t e r}} (\mathcal {I} _ {c} ^ {k}) := \sum_ {i = 1} ^ {K} \mathbb {1} (i = = k) \cdot \mathcal {F} _ {\mathrm {C o v 1}} ^ {(i)} \circ \mathcal {F} _ {\mathrm {C o v 2}} ^ {(i)} (\mathcal {I} _ {c} ^ {k}),
$$
where  $\mathbb{1}(\cdot)$  is the indicator function,  $\mathcal{I}_c^k$  is the conditioned image from task  $k$ , and  $\mathcal{F}_{\mathrm{Cov1}}^{(i)}, \mathcal{F}_{\mathrm{Cov2}}^{(i)}$  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.
MOE Hybrid Tasks Generalization
MOE Zero-Shot New Task Generalization
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_{\mathrm{task}}$ . As shown in Figure 2, our hyperNet first projects the task instruction  $c_{\mathrm{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 Uniconrol possesses the ability to transcend in-domain distributions for broader generalization. We demonstrate the zero-shot 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.
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.
Implementation. The UniControl is illustrated as Fig. 2 with Stable Diffusion, ControlNet, MOE Adapter, and Task-aware HyperNet consisting  $\sim 1.5\mathrm{B}$  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\times 10^{-5}$ . Our full-version UniControl model is trained on 16 Nvidia-A100 GPUs with the batch size of 4, requiring  $\sim 5$ , 000 GPU hours. We have also applied Safety-Checker as safeguards of results.
Figure 4: Visual comparison between official or re-implemented task-specific ControlNet and our proposed model. The example data is collected from our testing set sampled from COCO and Laion.
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 Laion-Aesthetics-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-to-one 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 image-condition-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
"Background: Photograph, Toronto Wet, Foreground: Some girls in colorful shirts standing by some pastries"
(a) Zero Shot Combination: Depth and Human Skeleton to Image
"Daniel Dae Kim smiles at the camera"
(c) Zero-Shot Task: Image Deblurring Example Result
Figure 5: (a)-(b): Example results of UniControl over hybrid (unseen combination) conditions with key words "background" and "foreground" attached in prompts. (c)-(e): Example results of UniControl on three unseen tasks (deblurring, colorization, inpainting).
"Background: Old City of Tzfat, Foreground: A set of twins who are taking a bath"
(b) Zero Shot Combination: Segmentation and Human Skeleton to Image
"Once Upon A Time: S03E19, 720p," kiss them goodbye
(d) Zero-Shot Task: Image Colorization Example Result
"Wedding hair and makeup with warm tones and red lips"
(e) Zero-Shot Task: Image Inpainting Example Result
Figure 6: User study between our method and official ControlNet checkpoints on six tasks. Our method outperforms ControlNet on all tasks.