UniControl

此仓库用于以下论文：

UniControl：一种用于野外可控视觉生成的统一扩散模型
秦灿 ^1,2, 张舒¹, 于宁¹, 冯一昊¹, 杨欣怡¹, 周英博¹, 王欢¹, 胡安·卡洛斯·尼布尔斯¹, 熊才明¹, 西尔维奥·萨瓦雷塞¹, 斯特凡诺·埃尔蒙³, 傅云², 徐然^1
1 Salesforce AI ² 东北大学 ³ 斯坦福大学
该工作完成于秦灿在Salesforce AI Research实习期间。

简介

我们提出了UniControl，这是一种新的生成式基础模型，它将广泛的条件到图像（C2I）可控任务整合到一个单一框架中，同时仍支持任意语言提示。UniControl能够实现像素级精确的图像生成，其中视觉条件主要影响生成的结构，而语言提示则指导风格和上下文。为了使UniControl具备处理多样化视觉条件的能力，我们对预训练的文本到图像扩散模型进行了增强，并引入了一个任务感知的HyperNet来调制这些扩散模型，从而使其能够同时适应不同的C2I任务。实验结果表明，UniControl通常优于同等规模的单任务控制方法。这种控制的多样性使UniControl成为可控视觉生成领域的一项重要进展。

更新

• 2023年5月18日：UniControl 论文已上传至arXiv。
• 2023年5月26日：UniControl推理代码和检查点向公众开放。
• 2023年5月28日：最新的UniControl模型检查点（1.4B参数，5.78GB）已更新。
• 2023年6月8日：最新的UniControl模型检查点已更新，现支持12项任务（Canny、HED、Sketch、Depth、Normal、Skeleton、Bbox、Seg、Outpainting、Inpainting、Deblurring 和 Colorization)！
• 2023年6月8日：训练数据集(MultiGen-20M)已完全公开。
• 2023年6月8日：训练代码已公开。:blush:
• 2023年7月6日：最新的UniControl模型v1.1 检查点已更新，现支持12项任务（Canny、HED、Sketch、Depth、Normal、Skeleton、Bbox、Seg、Outpainting、Inpainting、Deblurring 和 Colorization)！
• 2023年7月25日：HuggingFace演示API现已可用！
• 2023年7月25日：Safetensors模型现已可用！检查点。
• 2023年9月21日：UniControl已被NeurIPS 2023接收。:blush:

MultiGen-20M 数据集

这里有超过2000万张图像-提示-条件三元组在此，总大小***> 2TB***。它包含了所有12项任务（Canny、HED、Sketch、Depth、Normal、Skeleton、Bbox、Seg、Outpainting、Inpainting、Deblurring、Colorization），均已完全公开。

使用说明

环境准备

首先设置环境（需等待几分钟）。

conda env create -f environment.yaml
conda activate unicontrol

检查点准备

预训练的UniControl模型检查点保存在./ckpts/unicontrol.ckpt。

cd ckpts
wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol.ckpt 

你也可以使用最新训练的模型（ckpt和safetensors）：

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

如果你想从头开始训练，请按照ControlNet的步骤准备检查点初始化。ControlNet提供了一个简单的脚本，方便你完成这一操作。如果你的SD文件名为./ckpts/v1-5-pruned.ckpt，并希望脚本将处理后的模型（SD+ControlNet）保存到./ckpts/control_sd15_ini.ckpt，只需运行：

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

数据准备

请将训练数据集(MultiGen-20M)下载到./multigen20m。请执行：

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

然后解压所有文件。

模型训练（CUDA 11.0和Conda 4.12.0有效）

从零开始训练：

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

微调模型：

python train_unicontrol.py --ckpt ./ckpts/unicontrol.ckpt  --config ./models/cldm_v15_unicontrol.yaml --lr 1e-7

模型推理（CUDA 11.0 和 Conda 4.12.0 可用）

针对不同任务，请按以下方式运行代码。如果遇到显存不足（OOM）错误，请减少 --num_samples 参数的值。

如果您使用的是 safetensors 格式的模型，可以按照 ./load_model/load_safetensors_model.py 的方法加载模型。

Canny 边缘到图像生成：

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

HED 边缘到图像生成：

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

类似 HED 的草图到图像生成：

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

深度图到图像生成：

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

法线表面图到图像生成：

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

分割图到图像生成：

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

人体骨骼关键点到图像生成：

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

目标边界框到图像生成：

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

图像外扩填充：

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

图像修复填充：

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

图像去模糊：

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

图像上色：

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

Gradio 演示（应用演示视频，CUDA 11.0 和 Conda 4.12.0 可用）

我们为不同任务提供了 Gradio 演示工具。示例图片保存在 ./test_imgs 目录下。

对于所有任务（Canny、HED、Sketch、Depth、Normal、Human Pose、Seg、Bbox、Outpainting、Colorization、Deblurring、Inpainting），请运行以下代码：

python app/gradio_all_tasks.py

我们支持直接条件输入生成图像（如上所示）。如果您想直接上传条件图像，请在界面中取消勾选“条件提取”选项。

---

此外，我们还提供了特定任务的 Gradio 演示：

Canny 边缘到图像生成：

python app/gradio_canny2image.py

---

HED 边缘到图像生成：

python app/gradio_hed2image.py

---

类似 HED 的草图到图像生成：

python app/gradio_hedsketch2image.py

---

深度图到图像生成：

python app/gradio_depth2image.py

---

法线表面图到图像生成：

python app/gradio_normal2image.py

---

分割图到图像生成：

python app/gradio_seg2image.py

---

人体骨骼关键点到图像生成：

python app/gradio_pose2image.py

---

目标边界框到图像生成：

python app/gradio_bbox2image.py

---

图像外扩填充：

python app/gradio_outpainting.py

---

图像上色：

python app/gradio_colorization.py

---

图像去模糊：

python app/gradio_deblur.py

---

图像修复填充：

python app/gradio_inpainting.py

待办事项

• 数据准备
• 预训练任务推理
• Gradio 演示
• 模型训练
• HF Space
• Colab

小贴士

• 负面提示词有时非常有用：例如 monochrome, lowres, bad anatomy, worst quality, low quality 等。

• UniControl 在某些任务（如图像上色和去模糊）上，即使不使用任何文本提示词也能取得良好效果。

• 如果出现显存不足问题，可以尝试将 --num_samples 设置为 1。

引用

如果您觉得本项目对您的研究有所帮助，请引用我们的论文：

@article{qin2023unicontrol,
  title={UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild},
  author={Qin, Can and Zhang, Shu and Yu, Ning and Feng, Yihao and Yang, Xinyi and Zhou, Yingbo and Wang, Huan and Niebles, Juan Carlos and Xiong, Caiming and Savarese, Silvio and others},
  journal={arXiv preprint arXiv:2305.11147},
  year={2023}
}

致谢

本项目建立在 ControlNet 和 Stable Diffusion 这两座巨人的肩膀之上。衷心感谢它们！

Stable Diffusion https://github.com/CompVis/stable-diffusion

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

StyleGAN3 https://github.com/NVlabs/stylegan3

UniControl 快速上手指南

UniControl 是一个统一的扩散模型，能够在单一框架下支持多种可控图像生成任务（如边缘检测、深度图、分割、上色、去模糊等），同时保留自由文本提示的能力。

环境准备

系统要求：

• 操作系统：Linux (推荐) 或 macOS
• GPU：NVIDIA GPU (显存建议 8GB 以上，复杂任务需更大显存)
• CUDA 版本：11.0
• Conda 版本：4.12.0

前置依赖：

• Git
• wget
• Google Cloud SDK (gsutil) - 仅当需要下载训练数据集时必需

安装步骤

1. 克隆项目并配置环境

首先克隆仓库并创建 Conda 环境：

git clone https://github.com/salesforce/UniControl.git
cd UniControl
conda env create -f environment.yaml
conda activate unicontrol

2. 下载模型权重

进入检查点目录并下载预训练模型。你可以选择基础版或最新的 v1.1 版本（支持 12 种任务）：

mkdir -p ckpts
cd ckpts

# 选项 A: 下载最新 v1.1 模型 (推荐，支持更多任务)
wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol_v1.1.ckpt
# 或者使用 Safetensors 格式
wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol_v1.1.st

# 选项 B: 下载基础模型
# wget https://storage.googleapis.com/sfr-unicontrol-data-research/unicontrol.ckpt

注：如果下载速度慢，可尝试使用代理或寻找国内镜像源。

3. (可选) 准备自定义初始化权重

如果你需要从 Stable Diffusion 从头训练，需先运行脚本处理权重（基于 ControlNet 逻辑）：

# 假设你已有 v1-5-pruned.ckpt
python tool_add_control.py ./ckpts/v1-5-pruned.ckpt ./ckpts/control_sd15_ini.ckpt

基本使用

命令行推理示例

以下是最简单的使用方式，以 Canny 边缘转图像 为例。请确保当前处于 unicontrol 环境中，且 ckpts 目录下有对应的 .ckpt 文件。

python inference_demo.py --ckpt ./ckpts/unicontrol_v1.1.ckpt --task canny

支持的任务类型 (--task 参数)：

• canny: Canny 边缘检测
• hed: HED 边缘检测
• hedsketch: 素描草图
• depth: 深度图
• normal: 法线图
• seg: 语义分割图
• openpose: 人体骨骼姿态
• bbox: 物体边界框
• outpainting: 图像外绘 (Outpainting)
• inpainting: 图像修复 (Inpainting)
• blur: 图像去模糊
• grayscale: 图像上色

实用技巧：

• 如果遇到显存不足 (OOM) 错误，请在命令中添加 --num_samples 1。
• 部分任务（如上色 grayscale 和去模糊 blur）无需文本提示即可工作。
• 可添加负面提示词提升质量：monochrome, lowres, bad anatomy, worst quality, low quality。

启动本地 Web 界面 (Gradio Demo)

为了方便测试所有功能，可以启动本地的 Gradio 交互界面：

# 启动包含所有任务的综合界面
python app/gradio_all_tasks.py

启动后，在浏览器访问显示的本地地址（通常为 http://127.0.0.1:7860），即可上传条件图片并生成结果。如需针对特定任务启动独立界面，可运行 app/gradio_<task>2image.py 系列脚本。