UniControl: one diffusion model behind twelve kinds of control

The idea: one model, many control signals

UniControl is introduced as a generative foundation model that consolidates a wide array of controllable condition-to-image tasks within a singular framework, while still allowing for arbitrary language prompts. The division of labour it proposes is the interesting part. Visual conditions primarily influence the generated structures, and language prompts guide the style and context, so the model is not asked to make the prompt do everything a control map can do. To handle diverse visual conditions, the authors augment pretrained text-to-image diffusion models and introduce a task-aware HyperNet to modulate them, which is what allows adaptation to different tasks simultaneously rather than a separate model per task. The reported result is that UniControl often surpasses single-task controlled methods of comparable model sizes. The paper is UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild, from Salesforce AI, Northeastern University and Stanford, with the work done when Can Qin was an intern at Salesforce AI Research, and it was accepted to NeurIPS 2023.

The twelve tasks, and why the list is the point

The checkpoint supports twelve tasks, and the list is worth reading as a map of what people actually ask of control models. The edge and structure family covers Canny, HED, HED-like Sketch, Depth, Normal surface map and Skeleton. The layout family covers Bbox and Seg. The image-level family covers Outpainting, Inpainting, Deblurring and Colorization. Each of these is a different kind of constraint on the output, and a system that handles all of them with one set of weights handles a fundamentally different set of use cases from one that handles only depth. The training data, MultiGen-20M, is more than 20 million image-prompt-condition triplets with a total size of over 2TB, covering all twelve tasks and fully released, which is arguably the more valuable artefact for a research group than the checkpoint itself, because a dataset of aligned triplets is what you need to train a comparable model.

Setting up the environment and the weights

The setup is a conda environment created from a file in the repository, and the project notes it takes a few minutes.

conda env create -f environment.yaml
conda activate unicontrol

The pre-trained checkpoint is expected at ./ckpts/unicontrol.ckpt, and you fetch it into that directory.

The pre-trained checkpoint is expected at ./ckpts/unicontrol.ckpt, and you fetch it into that directory with a wget of the storage URL for unicontrol.ckpt, after changing into the ckpts folder. There is also a later version, v1.1, in both checkpoint and safetensors form.

wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol_v1.1.ckpt
wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol_v1.1.st

The update history in the project overview records the sizes and dates: the model checkpoint of 1.4B parameters and 5.78GB was updated on 05/28/23, the checkpoint supporting twelve tasks was updated on 06/08/23, the v1.1 checkpoint and the safetensors version arrived on 07/25/23, and a HuggingFace demo API was made available the same day.

Running inference, one flag per task

Inference goes through a single script with a task flag, which is the concrete payoff of the unified design. The Canny to image case looks like this.

python inference_demo.py --ckpt ./ckpts/unicontrol.ckpt --task canny

Every other condition is the same command with a different task name, for instance hed for the HED edge case and depth for the depth map case. The project warns that if you hit an out of memory error you should decrease the num_samples setting, which is a useful hint about the memory profile of running this model. If you use the safetensors model, the project says you can load it following the script in the load_model directory, which is the safer path and the one the project ends up recommending by publishing the safetensors weights.

Training from scratch means starting from ControlNet

If you want to train from scratch, the project tells you to follow ControlNet to prepare the checkpoint initialisation, because the model is built on a pretrained text-to-image diffusion model that has to be combined with the control branch before training starts. ControlNet provides a script for that, and the repository includes tool_add_control.py, plus a variant for the second generation of Stable Diffusion. The invocation takes the base model and the output location.

python tool_add_control.py ./ckpts/v1-5-pruned.ckpt ./ckpts/control_sd15_ini.ckpt

The naming in that command is instructive, since the prepared checkpoint is written to control_sd15_ini.ckpt, indicating a Stable Diffusion 1.5 base. Once that exists, training from scratch uses the initialised checkpoint, a config file and a learning rate.

python train_unicontrol.py --ckpt ./ckpts/control_sd15_ini.ckpt --config ./models/cldm_v15_unicontrol_v11.yaml --lr 1e-5

Finetuning the released checkpoint instead uses the released weights, a different config file, cldm_v15_unicontrol.yaml, and a much smaller learning rate of 1e-7, which is the sane starting point for most people.

The project notes that CUDA 11.0 and Conda 4.12.0 are a working combination, which dates the environment and is worth taking seriously before you start installing.

Getting the training data, which is the real download

The data step is where the disk requirement comes from. You are told to download the MultiGen-20M training dataset into a directory named multigen20m, and the transfer uses Google Cloud Storage tooling with the parallel flag.

cd multigen20m
gsutil -m cp -r gs://sfr-unicontrol-data-research/dataset ./

Then you unzip all the files. With a dataset over 2TB, this step dominates the time and space budget of the whole exercise, and it is worth being honest about that before you plan around the model. The payoff is a fully released dataset of more than 20 million triplets where the image, the prompt and the condition are aligned, which is the kind of data that is expensive to assemble and expensive to verify. The repository also has a dataset_eval.py at the top level, plus data/ and annotator/ directories, which suggests the annotation and evaluation tooling is included alongside the data.

What the repository contains, and its current state

The top level of the repository shows the structure of a research codebase built on an existing diffusion stack. There is a cldm/ directory, which is the latent diffusion model code the project builds on, an ldm/ directory, a train_util/ directory and a utils.py, along with config.py, share.py and inference_demo.py at the top level. The models/ directory holds the yaml configs referenced in the training commands, load_model/ holds the safetensors loading script, and there are directories for figures, fonts, test images and a data directory. The licensing and governance files are the usual set for a Salesforce repository, with an AI_ETHICS.md, a SECURITY.md, a CODEOWNERS file, a CODE_OF_CONDUCT.md and a CONTRIBUTING.md, and the licence is Apache-2.0. The last push to the repository was on 2026-06-02, and the project publishes no GitHub releases, so the checkpoint and the code are what you get and there is no packaged release to install.

Editorial conclusion

UniControl fits a researcher who wants to compare control conditions across tasks with one model rather than twelve, since the same checkpoint handles all twelve tasks and the training data and code are both released, and it fits anyone who wants a controllable generation baseline with published numbers rather than a demo. It does not fit a team expecting a maintained product, because the last push to the repository was on 2026-06-02 and the update history in the project overview stops in 2023, with the paper accepted to NeurIPS 2023. It also does not fit a team without a large GPU, since the checkpoint is 1.4B parameters and 5.78GB and the training data is over 2TB. Before you start, read the environment file rather than assuming a current PyTorch stack, and note that the training-from-scratch path depends on ControlNet for checkpoint initialisation, which is a separate dependency to account for.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. README
5. salesforce/UniControl on GitHub