SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shihao Zhao†
The University of Hong Kong shzhao@cs.hku.hk
Dongdong Chen*
cddlyf@gmail.com
Yen-Chun Chen
yen-chun.chen@microsoft.com
Jianmin Bao
jianmin.bao@microsoft.com
Shaozhe Hao
szhao@cs.hku.hk
Lu Yuan
luyuan@microsoft.com
Kwan-Yee K. Wong*
kykwong@cs.hku.hk
Text-to-Image diffusion models have made tremendous progress over the past two years, enabling the generation of highly realistic images based on open-domain text descriptions. However, despite their success, text descriptions often struggle to adequately convey detailed controls, even when composed of long and complex texts. Moreover, recent studies have also shown that these models face challenges in understanding such complex texts and generating the corresponding images. Therefore, there is a growing need to enable more control modes beyond text description. In this paper, we introduce Uni-ControlNet, a unified framework that allows for the simultaneous utilization of different local controls (e.g., edge maps, depth map, segmentation masks) and global controls (e.g., CLIP image embeddings) in a flexible and composable manner within one single model. Unlike existing methods, Uni-ControlNet only requires the fine-tuning of two additional adapters upon frozen pre-trained text-to-image diffusion models, eliminating the huge cost of training from scratch. Moreover, thanks to some dedicated adapter designs, Uni-ControlNet only necessitates a constant number (i.e., 2) of adapters, regardless of the number of local or global controls used. This not only reduces the fine-tuning costs and model size, making it more suitable for real-world deployment, but also facilitate composability of different conditions. Through both quantitative and qualitative comparisons, Uni-ControlNet demonstrates its superiority over existing methods in terms of controllability, generation quality and composability. Code is available at https://github.com/ShihaoZhaoZSH/Uni-ControlNet.
In recent two years, diffusion models [1-10] have gained significant attention due to their remarkable performance in image synthesis tasks. Therefore, text-to-image (T2I) diffusion models [6, 7, 11-19] have emerged as a popular choice for synthesizing high-quality images based on textual inputs. By training on large-scale datasets with large models, these T2I diffusion models demonstrate exceptional ability in creating images that closely resemble the content described in text descriptions, and facilitate
*Corresponding Author, † Intern at Microsoft
the connection between textual and visual domains. The substantially improved generation quality in capturing intricate texture details and complex relationships between objects, makes them highly suitable for various real-world applications, including but not limited to content creation, fashion design, and interior decoration.
However, text descriptions often prove to be either inefficient or insufficient to accurately convey detailed controls upon the final generation results, e.g., control the fine-grained semantic layout of multiple objects, not to mention the challenge in understanding complex text descriptions for such models. As a result, there is a growing need to incorporate more additional control modes (e.g., user-drawn sketch, semantic mask) alongside the text description into such T2I diffusion models. This necessity has sparked considerable interest from both academia and industry, as it broadens the scope of T2I generation from a singular function to a comprehensive system.
Very recently, there are some attempts [20-22] studying controllable T2I diffusion models. One representative work,Composer [20], explores the integration of multiple different control signals together with the text descriptions and train the model from scratch on billion-scale datasets. While the results are promising, it requires massive GPU resources and incurs huge training cost, making it unaffordable for many researchers in this field. Considering there are powerful pretrained T2I diffusion models (e.g., Stable Diffusion [6]) publicly available, ControlNet [21], GLIGEN [23] and T2I-Adapter [22] directly incorporate lightweight adapters (or extra modules) into frozen T2I diffusion models to enable additional condition signals. This makes fine-tuning more affordable. However, one drawback is that they need one independent adapter for each single condition, resulting in a linear increase in fine-tuning cost and model size along as the number of the control conditions grows, even though many conditions share similar characteristics. Additionally, this also makes composability among different conditions remains a formidable challenge.
In this paper, we propose Uni-ControlNet, a new framework that leverages lightweight adapters to enable precise controls over pre-trained T2I diffusion models. As shown in Table 1, Uni-ControlNet can not only handle different conditions within one single model but also supports composable control. By contrast, the existing methods fail to achieve this unified framework within one single model. Besides, even for those methods that support composite control, they perform poorly in terms of composability as illustrated in Section 4.
Unlike previous methods, Uni-ControlNet categorizes various conditions into two distinct groups: local conditions and global conditions. Accordingly, we only add two additional adapters, regardless of the number of local and global controls involved. This design choice not only significantly reduces both the whole fine-tuning cost and the model size, making it highly efficient for deployment, but also facilitates the composability of different conditions. To achieve this, we dedicatedly design the adapters for local and global controls. Specifically, for local controls, we introduce a multi-scale condition injection strategy that uses a shared local condition encoder adapter. This adapter first converts the local control signals into modulation signals, which are then used to modulate the incoming noise features. And for global controls, we employ another shared global condition encoder to convert them into conditional tokens, which are concatenated with text tokens to form the extended prompt and interacted with the incoming features via cross-attention mechanism. Interestingly, we find these two adapters can be separately trained without the need of additional joint training, while still supporting the composition of multiple control signals. This finding adds to the flexibility and ease of use provided by Uni-ControlNet.
By only training on 10 million text-image pairs with 1 epoch, our Uni-ControlNet demonstrates highly promising results in terms of fidelity and controllability. Figure 1 provides visual examples showcasing the effectiveness of Uni-ControlNet when using either one or multiple conditions. To gain further insights, we perform in-depth ablation analysis and compare our newly proposed adapter designs with those of ControlNet [21], GLIGEN [23] and T2I-Adapter [22]. The analysis results reveal the superiority of our adapter designs, emphasizing their enhanced performance over the counterparts offered by ControlNet, GLIGEN and T2I-Adapter.
Text-to-Image Generation is an emerging field that aims to generate realistic images from text descriptions. To address this challenging task, various approaches have been proposed in the past years. Early works [24-26] primarily adopted Generative Adversarial Networks (GANs) and were
Figure 1: Visual results of our proposed Uni-ControlNet. The top and bottom two rows are results for single condition and multi-conditions respectively.
Table 1: Comparisons of different controllable diffusion models.  $N$  is the number of conditions. We define the fine-tuning cost as the number of times the model needs to be fine-tuned on  $N$  conditions. As Composer is trained from scratch, both fine-tuning cost and adapter number are not applicable. For T2I-Adapter,  $(+1)$  indicates that further joint fine-tuning is required on the  $N$ -based adapters along with an additional fuser to achieve composable conditions.
Fine-tuning	Composable Control	Fine-tuning Cost	Adapter Number
Composer	X	✓	-	-
ControlNet	✓	✓	N	N
GLIGEN	✓	X	N	N
T2I-Adapter	✓	✓	N(+1)	N(+1)
Uni-ControlNet (Ours)	✓	✓	2	2
often trained on specific domains. However, they faced two key challenges, i.e., training instability and poor generalization ability to open-domain scenarios. Motivated by the success of GPT models [27-30], recent works [31-34] have explored the use of autoregressive models for text-to-image generation and train on web-scale image-text pairs, which start to show strong generation capability under the zero-shot setting for open-domain scenarios. Another approach is the diffusion models [6, 7, 11, 35-41], originally proposed by [1, 2]. Diffusion models comprise a forward process that gradually adds noise to natural images and a backward process that learns to denoise them back to generate clean output. They demonstrate stronger capability in modeling fine-grained structures and texture details compared to autoregressive models. Recently, vast variants of diffusion models have been developed, such as DALLE-2 [11], which uses one prior model and one decoder model to generate images from CLIP latent embeddings. Another phenomenal T2I diffusion model is Stable Diffusion (SD), which scaled up the latent diffusion model [6] with larger model and data scales, and made the pre-trained models publicly available. In this paper, we use SD as a base model and explore how to enable more control signals beyond the text description for pre-trained T2I diffusion models in an efficient and composable way.
Controllable Diffusion Models are designed to enable T2I diffusion models to accept more user controls for guiding the generation results. They have garnered increasing attention very recently.
Figure 2: The overall framework of our proposed Uni-ControlNet.
Broadly speaking, there are two strategies for implementing controllable diffusion models: training from scratch [20] and fine-tuning lightweight adapters [21, 22] on frozen pretrained T2I diffusion models. In the case of training from scratch,Composer [20] trains one big diffusion model from scratch to achieve great controllability for both single and multi-conditions. It obtains remarkable generation quality but comes with huge training cost. In contrast, ControlNet [21], GLIGEN [23] and T2I-Adapter [22] propose to introduce lightweight adapters (or extra modules) into publicly available SD models. By only fine-tuning the adapters while keeping original SD models frozen, they significantly reduce the training cost and make it affordable for the research community. However, all the ControlNet, GLIGEN and T2I-Adapter utilize independent adapters for each condition, resulting in increased fine-tuning cost and model size when handling increased number of conditions. Moreover, GLIGEN does not support composite control over different conditions. And different adapters in Multi-ControlNet [21], a version of ControlNet that allow composite control, are isolated from one another, limiting their composability. By testing CoAdapter [22], which is jointly trained using different T2I-Adapters, we find that it also exhibits inadequate performance in generating composable conditions. Our proposed Uni-ControlNet follows the second line of fine-tuning adapters and is much less expensive thanComposer, while addressing the above limitations of ControlNet, GLIGEN and T2I-Adapter. It groups conditions into two groups, i.e., local controls and global controls, and only requires two additional adapters accordingly. Thanks to our newly designed adapter structure, Uni-ControlNet is not only efficient in terms of training cost and model sizes, but also surpasses ControlNet, GLIGEN and T2I-Adapter in controllability and quality.
A typical diffusion model involves two processes: a forward process which gradually adds small amounts of Gaussian noise onto the sample in  $T$  steps, and a corresponding backward process containing learnable parameters to recover input images by estimating and eliminating the noise. In this paper, we use SD as our example base model to illustrate how to enable diverse controls with our Uni-ControlNet. SD incorporates the UNet-like structure [42] as its denoising model, which consists of an encoder, a middle block, and a decoder, with 12 corresponding blocks in each of the encoder and decoder modules. For brevity, we denote the encoder as  $F$ , the middle block as  $M$ , and the decoder as  $G$ , with  $f_{i}$  and  $g_{i}$  denoting the output of the  $i$ -th block in the encoder and decoder, and  $m$  denoting the output of the middle block, respectively. It is important to note that, due to the adoption of skip connections in UNet, the input for the  $i$ -th block in the decoder is given by:
$$
\left\{ \begin{array}{l} \operatorname {c o n c a t} \left(m, f _ {j}\right) \\ \operatorname {c o n c a t} \left(g _ {i - 1}, f _ {j}\right) \end{array} \quad \text {w h e r e} \quad i = 1, \quad i + j = 1 3. \right. \tag {1}
$$
Skip connections allow the decoder to directly utilize features from the encoder and thereby help minimize the information loss. In SD, cross-attention layers are employed to capture semantic information from the input text description. Here we use  $Z$  to denote the incoming noise features and  $y$  to denote text token embeddings encoded by the language encoder. The  $Q, K, V$  in cross-attention can be expressed as:
$$
Q = W _ {q} (Z), K = W _ {k} (y), V = W _ {v} (y), \tag {2}
$$
where  $W_{q}, W_{k}$  and  $W_{v}$  are projection matrices.
Figure 3: Details of the local and global control adapters.
3.2 Control Adapter
In this paper, we consider seven example local conditions, including Canny edge [43], MLSD edge [44], HED boundary [45], sketch [46, 47], Openpose [48], Midas depth [49], and segmentation mask [50]. We also consider one example global condition, i.e., global image embedding of one reference content image that is extracted from the CLIP image encoder [51]. This global condition goes beyond simple image features and provides a more nuanced understanding of the semantic content of the condition image. By employing both local and global conditions, we aim to provide a comprehensive control over the generation process. We show the overview of our pipeline in Figure 2, and the details of local control adapter and global control adapter are given in Figure 3.
Local Control Adapter: For our local control adapter, we have taken inspiration from ControlNet. Specifically, we fix the weights of SD and copy the structures and weights of the encoder and middle block, designated as  $F'$  and  $M'$  respectively. Thereafter, we incorporate the information from the local control adapter during the decoding process. To achieve it, we ensure that all other elements remain unchanged while modifying the input of the  $i$ -th block of the decoder as
$$
\left\{ \begin{array}{l} \operatorname {c o n c a t} \left(m + m ^ {\prime}, f _ {j} + z e r o \left(f _ {j} ^ {\prime}\right)\right) \\ \operatorname {c o n c a t} \left(g _ {i - 1}, f _ {j} + z e r o \left(f _ {j} ^ {\prime}\right)\right) \end{array} \quad \text {w h e r e} \quad i = 1, \quad i + j = 1 3. \right. \tag {3}
$$
where zero represents one zero convolutional layer whose weights increase from zero to gradually integrate control information into the main SD model. In contrast to ControlNet that adds the conditions directly to the input noise and sends them to the copied encoder, we opt for a multiscale condition injection strategy. Our approach involves injecting the condition information at all resolutions. In detail, we first concatenate different local conditions along the channel dimension and then use a feature extractor  $H$  (stacked convolutional layers) to extract condition features at different resolutions. Subsequently, we select the first block of each resolution (i.e.,  $64 \times 64$ ,  $32 \times 32$ ,  $16 \times 16$ ,  $8 \times 8$ ) in the copied encoder (i.e., the Copied Encoder in Figure 3) for condition injection. For the injection module, we take the inspiration from SPADE [52] and implement Feature Denormalization (FDN) that uses the condition features to modulate the normalized (i.e.,  $\text{norm}(\cdot)$ ) input noise features:
$$
F D N _ {r} \left(Z _ {r}, c _ {l}\right) = \operatorname {n o r m} \left(Z _ {r}\right) \cdot \left(1 + \operatorname {c o n v} _ {\gamma} \left(\operatorname {z e r o} \left(h _ {r} \left(c _ {l}\right)\right)\right) + \operatorname {c o n v} _ {\beta} \left(\operatorname {z e r o} \left(h _ {r} \left(c _ {l}\right)\right)\right), \right. \tag {4}
$$
where  $Z_{r}$  denotes noise features at resolution  $r$ ,  $c_{l}$  is the concatenated local conditions,  $h_{r}$  represents the output of the feature extractor  $H$  at resolution  $r$ , and  $conv_{\gamma}$  and  $conv_{\beta}$  refer to learnable convolutional layers that convert condition features into spatial-sensitive scale and shift modulation coefficients. We will ablate different local feature injection strategies in following sections.
Global Control Adapter: For global controls, we use the image embedding of one condition image extracted from CLIP image encoder as the example. Inspired by the fact that the text description in T2I diffusion models can be also viewed as one kind of global control without explicit spatial guidance, we project the global control signals into condition embeddings by using a condition encoder  $h_{g}$ . The condition encoder consists of stacked feedforward layers, which aligns the global control signals with the text embeddings in SD. Next, we reshape the projected condition embeddings into  $K$  global tokens ( $K = 4$  by default) and concatenate them with the original  $K_{0}$  text tokens
Figure 4: More visual results of Uni-ControlNet. The top two rows show results of a single condition, with columns 1-7 for local conditions and columns 8-9 for global condition. 3rd row shows the results of combining two local conditions, while row 4-th shows the results of integrating a local condition with a global condition. There is no text prompt for the examples in 4-th row.
to create an extended prompt  $y_{ext}$  (total token number is  $K + K_0$ ) which serves as the input to all cross-attention layers in both main SD model and control adapters:
$$
y _ {e x t} = \left[ y _ {1} ^ {t}, y _ {2} ^ {t}, \dots , y _ {K _ {0}} ^ {t}, \lambda * y _ {1} ^ {g}, \lambda * y _ {2} ^ {g}, \dots , \lambda * y _ {K} ^ {g} \right], \text {w h e r e} y _ {i} ^ {g} = h _ {g} \left(c _ {g}\right) [ (i - 1) \cdot d \sim i \cdot d ], i \in [ 1, K ] \tag {5}
$$
where  $y^{t}$  and  $y^{g}$  represent the original text tokens and global condition tokens respectively, and  $\lambda$  is a hyper-parameter that controls the weight of the global condition.  $c_{g}$  denotes the global condition and  $d$  is the dimension of text token embedding.  $h_{g}(\cdot)[i_{s}\sim i_{e}]$  represents the sub-tensor of  $h_{g}(\cdot)$  that contains elements from the  $i_{s}$ -th to the  $i_{e}$ -th positions. Finally, the  $Q,K,V$  cross-attention operation in all cross-attention layers is changed to:
$$
Q = W _ {q} (Z), K = W _ {k} \left(y _ {\text {e x t}}\right), V = W _ {v} \left(y _ {\text {e x t}}\right), \tag {6}
$$
As the local control signals and global control signals often contain different amounts of condition information, we empirically find that directly joint fine-tuning these two types of adapters will produce poor controllable generation performance. Therefore, we opt to fine-tune these two types of adapters separately so that both of them can be sufficiently trained and contribute effectively to the final generation results. When fine-tuning each adapter, we employ a predefined probability to randomly dropout each condition, along with an additional probability to deliberately keep or drop all conditions. For the dropped conditions, we set the value of the corresponding input channels to 0. This can facilitate the model to learn generating the results based on one or multiple conditions simultaneously. Interestingly, by directly integrating these two separately trained adapters during inference, our Uni-ControlNet can already well combine global and local conditions together in a composable way, without the need of further joint fine-tuning. In Section 4.3, we will provide more detailed analysis about different fine-tuning strategies.
Implementation Details. To fine-tune our model, we randomly sample 10 million text-image pairs from the LAION dataset [53] and fine-tune Uni-ControlNet for 1 epoch. We use the AdamW optimizer [54] with a learning rate of  $1 \times 10^{-5}$  and resize the input images and local condition maps to  $512 \times 512$ . As described, the local and global control adapters are fine-tuned separately by
default. During inference, we merge the two adapters and adopt DDIM [55] for sampling, with the number of time steps set to 50 and the classifier free guidance scale [56] set to 7.5. During training, the hyper-parameter  $\lambda$  in Equation 6 is with a fixed value 1. At inference time, when there is no text prompt,  $\lambda$  remains at 1, while when there is a text prompt, the value is adjusted to around 0.75, depending on the intended weight between the text and global condition. As explained in Section 3.2, we employ 7 local conditions (Canny edge, MLSD edge, HED boundary, sketch, Openpose, Midas depth, and segmentation mask) and 1 global control condition (CLIP image embeddings) for control. As annotating a sketch dataset can be challenging, in our experiment, we initially obtain the HED boundary [45] of an image and subsequently utilize a sketch simplification method [46, 47] to generate the sketch for the training sample. Regarding the pose condition, as not all images in the dataset include humans, we opt to not drop the pose condition during training to ensure that the pose condition is fully trained. Detailed structures of global and local condition adapters can be found in the appendix.
4.1 Controllable Generation Results
In Figure 4, we provide more controllable generation results of Uni-ControlNet in both single and multi-condition setups. Notably, for visualization purposes, we use the original condition images to denote their CLIP image embeddings. It can be seen that our Uni-ControlNet can produce very promising results in terms of both controllability and generation fidelity. For example, in the case of a single sketch condition with the text prompt "Dog, wild" (rows 1-2, column 4), the resulting image accurately depicts a vivid dog and a background of grass and trees that align well with the given sketch condition. Similarly, when presented with the global CLIP image embedding conditions with the prompt "Golden retriever" (rows 1-2, columns 8-9), our model can seamlessly change the background of the dog from the wild to a room. Moreover, our model also handles multi-condition settings well, as demonstrated in the example of "A man on the mountains" (row 3, columns 7-9), where the combination of a sketch and a pose produces a cohesive and detailed image of a man on a mountainside. When presented with a local depth map and global CLIP image embeddings without any prompt (row 4, columns 1-3), our model produces an image of a forest, taking the contour of an elephant, which harmonizes with both the depth map and the content of the source global image.
Here we compare our Uni-ControlNet with ControlNet (Multi-ControlNet) [21], GLIGEN [23] and T2I-Adapter (CoAdapter) [22]. Since Composer [20] is not open-sourced and trained from scratch, we do not include it in comparisons.