Introduction

The full title of the ControlNet paper is Adding Conditional Control to Text-to-Image Diffusion Models. The study known as ControlNet discusses how to enable efficient knowledge transfer from a pretrained large diffusion model so that it follows an input condition. Although its title mentions diffusion models, the direction of this work is completely different from that of the foundational diffusion papers reviewed earlier. It is more naturally connected to parameter-efficient fine-tuning, a recent line of research for training large language models (LLMs). In my own experiments with the code, the ControlNet architecture proposed by the authors could be trained comfortably on a consumer GPU rather than a server-class GPU. Its most striking advantage is that ControlNet can be trained for virtually any input condition. The core idea of the method is called a zero convolution. Let us work through how it makes input conditions so flexible.

Input Conditions in Diffusion Models

Methods for providing input conditions to diffusion models already existed. I have not yet covered the latent diffusion paper, which underpins today's popular Stable Diffusion, so here is a brief introduction.

Suppose we have trained a vector-quantized encoder/decoder that retains only meaningful semantic information while filtering out high-frequency features that contribute little to image generation. In other words, it reduces an image $x$ to a lower-resolution latent image. Learning to recover this compressed latent through a diffusion process is the training and sampling process of DDPM or DDIM as we usually understand it.

In an earlier post, we examined a method that supplies time and class-label information to attention pooling as projection embeddings (review of diffusion-process conditioning). Extending that idea, if we have a task-specific encoder $\tau_\theta$ capable of projecting a condition for a particular purpose into an embedding, then during diffusion-model training we can condition the model through an attention layer using the condition vector extracted by $\tau_\theta$. For example, consider the following image-description pair from the BLIP paper:

Description: The car is driving past a small old building

If the embedding extracted by an arbitrary text encoder, such as CLIP's text encoder, is supplied as a condition during image generation (the reverse process), the diffusion model can perform text-to-image generation simply by adding the prompt condition to its prior during sampling.

This approach can do more than produce high-quality images from prompts. Given trained encoders for different modalities, it can also train a diffusion process through attention-pooling conditioning.

Drawbacks of Naive Conditioning

The problems with these methods are fairly clear.

First, for a diffusion model to learn a particular condition, the score network must reflect that condition effectively in the generated image, which requires a very large training dataset. Vision-language (VL) tasks have been studied extensively in multimodal research, so massive datasets for models such as CLIP and ALIGN have been assembled. Other modalities—pose-to-image, semantic-map-to-image, and so on—do not enjoy the same abundance. Compared with datasets such as LAION-5B, which formed the training base for Stable Diffusion, it is difficult to build large, specialized datasets for properties such as object shape or pose because of various practical limitations. The difference can be on the order of tens of thousands of times.

Second, deriving solutions for processes such as image generation and manipulation from enormous datasets is extremely resource-intensive. Even if we overcome the disparity in dataset size, training a pretrained network remains a substantial barrier.

Finally, image-processing problems can take such diverse—and still evolving—forms that their boundaries are impossible to predict in advance. If we rely only on greedy solutions to an open-ended problem, such as constraining the diffusion process or altering attention activations, it becomes difficult to optimize tasks that demand high-level understanding, including depth and pose. Put in terms of latent diffusion, the latent diffusion process depends on the embedding output of a pretrained task-specific encoder for conditioning: constraints are imposed merely by injecting that embedding during sampling. This cannot be optimal when learning a broad variety of tasks. We therefore need a method that supports end-to-end training. As the figure below illustrates, conventional approaches separate the conditioning component from the actual diffusion-model training instead of training them end to end.

ControlNet: An End-to-End Neural Network

To overcome these limitations of conventional conditioning, the authors propose a new transfer-learning architecture. ControlNet can be summarized as follows:

• It copies the diffusion model's parameters and places the new training framework in parallel with the original parameters. The two branches are called the trainable copy and the locked copy.
• The locked copy retains the representations required for the original network's image-generation capability.
• The trainable copy is trained on various task-specific datasets to learn conditional control.
• The locked and trainable copies are connected through zero convolutions, which are themselves trainable layers.

That description is easier to understand with a diagram.

When this input-$x$-to-output-$y$ structure is incorporated into a diffusion process, it corresponds to taking the noisy latent vector $z_t$ at one time step as input and predicting the noisy latent vector $z_{t-1}$ at the next step. The gray neural network is the original diffusion model. By freezing its parameters, we can leave the pretrained model untouched without damaging its image-generation ability.

Think of the frozen branch on the left as remaining still while only the active branch on the right learns the condition. Because the latter is a trainable copy, the process is a form of fine-tuning designed to disturb the original parameters as little as possible.

Method

Let us now examine mathematically how this architecture performs conditioning effectively. A neural-network block to be conditioned might be a bottleneck block in a ResNet or a multi-head attention block in a Transformer.

Consider a two-dimensional feature such as an image. Given a feature map $x \in \mathbb{R}^{h \times w \times c}$, a neural-network block $\mathcal{F}_\Theta(\cdot)$ transforms the input feature map $x$ through the parameters $\Theta$ in that block:

$$y = \mathcal{F}_\Theta(x)$$

This is the operation shown in part (a) of the preceding figure. We now lock the parameters $\Theta$ so that they are not trained. An identical trainable copy, with parameters $\Theta_c$, receives the input condition $c$ and learns from it, unlike the locked branch.

The added branch does not require twice as much memory during training because the network does not need to retain activations for the locked branch. Gradients computed through backpropagation optimize only the trainable model.

Zero Convolution

The most important concept in this paper, the zero convolution, is used when the branches are added together. One such layer is attached before and another after each neural block. Let us call them $\mathcal{Z}_{\Theta_1}(\cdot)$ and $\mathcal{Z}_{\Theta_2}(\cdot)$. A zero convolution must not change the feature-map dimensions, so it is a $1 \times 1$ convolution whose weights and biases are both initialized to zero.

Following the diagram above, adding the output of the conditioning function to the original output $y$ gives:

$$y_c = \mathcal{F}_\Theta(x) + \mathcal{Z}_{\Theta_2}(\mathcal{F}_{\Theta_c}(x + \mathcal{Z}_{\Theta_1}(c)))$$

We can now see why a zero convolution, whose weights and biases begin at $0$, is used. Because both are zero, it initially produces an output of zero regardless of its input:

$$\begin{cases} \mathcal{Z}_{\Theta_1}(c) = 0 \\ \mathcal{F}_{\Theta_c}(x+\mathcal{Z}_{\Theta_1}(c)) = \mathcal{F}_{\Theta_c}(x) = \mathcal{F}_{\Theta}(x) \\ \mathcal{Z}_{\Theta_2}(\mathcal{F}_{\Theta_c}(x + \mathcal{Z}_{\Theta_1}(c))) = \mathcal{Z}_{\Theta_2}(\mathcal{F}_{\Theta_c}(x)) = 0 \end{cases}$$

Thus training begins with $y_c = y$. This means that when training starts, the input-output relationship of the ControlNet architecture is identical to that of the pretrained diffusion model. The added network depth therefore has no effect before optimization begins.

Gradient Flow in a Zero Convolution

We can derive the local gradients of a zero convolution with a $1 \times 1$ convolutional structure. Given an input feature map $I \in \mathbb{R}^{h \times w \times c}$, the forward pass is:

$$\mathcal{Z}(I; \{W, B\})_{p, i} = B_i + \sum_{j=1}^{c} I_{p, j}W_{i, j}$$

Before optimization, a zero convolution has $W = 0$ and $B = 0$. At every point where $I_{p,j}$ is nonzero, we therefore obtain:

$$\begin{cases} \frac{\partial \mathcal{Z}(I; \{W, B\})_{p, i}}{\partial B_i} = 1 \\ \frac{\partial \mathcal{Z}(I; \{W, B\})_{p, i}}{\partial I_{p, j}} = W_{i,j} = 0 \\ \frac{\partial \mathcal{Z}(I; \{W, B\})_{p, i}}{\partial W_{i, j}} = I_{p, j} \neq 0 \end{cases}$$

Although the gradient with respect to the input is zero, the gradients with respect to the weights and biases are not, so the layer can learn. After just the first step, using $\odot$ to denote the Hadamard product,

$$W^\ast = W-\beta_\text{lr} \cdot \frac{\partial \mathcal{L}}{\partial \mathcal{Z}(I; \{W, B\})} \odot \frac{\partial \mathcal{Z}(I; \{W, B\})}{\partial W} \neq 0,$$

the weights become nonzero. Consequently, at the very next step,

$$\frac{\partial \mathcal{Z}(I; \{W^\ast, B\})_{p, i}}{\partial I_{p, j}} = W^\ast_{i,j} \neq 0,$$

and training proceeds normally.

Stable Diffusion + ControlNet

The figure above shows the architecture obtained by implementing the structure just described in Stable Diffusion. Its loss is the ordinary diffusion objective augmented with a task-specific condition $c_f$:

$$\mathcal{L} = \mathbb{E}_{z_0, t, c_t, c_f, \epsilon \sim \mathcal{N}(0, 1)}\left( \lVert \epsilon - \epsilon_\theta(z_t, t, c_t, c_f) \rVert_2^2 \right)$$

Results

Canny Edge

Hough Line

Scribble

HED Edge

Pose

Segmentation

Depth

Cartoon Line Drawing

Trying the Official Code

The official GitHub repository contains the source code. It will generally run on a local machine, though a server is preferable for reliable operation. I will include the procedure here, but consider visiting the original repository and leaving a star. First, clone it:

git clone https://github.com/lllyasviel/ControlNet.git

Then create the Conda environment:

cd ControlNet
conda env create -f environment.yaml
conda activate control

Next, download the desired model and Stable Diffusion checkpoint from the Hugging Face page. Put the Stable Diffusion checkpoint under ControlNet/models and the detectors under ControlNet/annotator/ckpts.

Downloading All Detectors (Run from the ControlNet Repository)

cd ./annotator/ckpts
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_tiny_512_fp32.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/network-bsds500.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/upernet_global_small.pth

If you do not need every detector, run curl only for the files you want.

Downloading All Models (Run from the ControlNet Repository)

cd ./models
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_canny.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_depth.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_hed.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_mlsd.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_normal.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_openpose.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_scribble.pth
curl -LO https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_seg.pth

Again, if you do not want them all, run curl only for the files you need.

Running the Demo API

Running the desired model is as simple as:

python gradio_something2something.py

If you launch it on a server but want to connect from your local machine, make one small change. At the bottom of each gradio_something2something.py file, you will find:

block.launch(server_name='0.0.0.0')

Change it to:

block.launch(server_name='0.0.0.0', share=True)

I tried running:

python gradio_scribble2image_interactive.py

If everything works, you will see something like this:

The message roughly says that the public URL remains valid for 72 hours. I was connected to the server through Yonsei University's VPN, but opened the link on my laptop.

The following screen appears. Let us check whether it actually works.

My humble drawing skills... Hang in there, ControlNet.

Pressing Run starts the DDIM sampler.

This paper offered the reassuring hope that being bad at drawing is not a major obstacle in life. That is all there is to it. A GPU with at least about 10 GB of memory appears to be enough.