[NeurIPS 2023] Uni-ControlNet：面向文本到图像扩散模型的一体化控制方法

Uni-ControlNet 的官方实现：一种面向文本到图像扩散模型的一体化控制方法，已被 NeurIPS 2023 接收。

项目页面 | 论文（ArXiv）

⏳ ：待办事项

• 发布训练代码
• 发布测试代码
• 发布预训练模型

💡 ：方法

Uni-ControlNet 是一种新颖的可控扩散模型，能够在单个模型中以灵活且可组合的方式同时利用多种局部控制和全局控制。这是通过引入两个适配器——局部控制适配器和全局控制适配器来实现的，无论使用多少种局部或全局控制信号均可适用。这两个适配器可以分别进行训练，无需联合训练，同时仍支持多组控制信号的组合。

以下是不同可控扩散模型的对比。N 表示条件的数量。Uni-ControlNet 不仅在控制条件数量增加时能够降低微调成本和模型规模，还能够方便地实现不同条件的组合。

⚙ ：环境搭建

首先创建一个新的 conda 环境：

conda env create -f environment.yaml
conda activate unicontrol

然后下载 预训练模型 (或在此处) 并将其放置在 ./ckpt/ 文件夹中。该模型基于 Stable Diffusion v1.5 构建。

💻 ：测试

您可以通过以下命令启动 Gradio 演示：

python src/test/test.py

此命令将加载下载的预训练权重并启动应用程序。我们包含了七种示例局部条件：Canny 边缘、MLSD 边缘、HED 边界、草图、Openpose、Midas 深度图、分割掩码，以及一个示例全局条件：内容。

您可以先上传一张源图像，我们的代码会自动检测其草图。随后，Uni-ControlNet 将根据草图和文本提示生成样本，在本例中提示为“机器人蜘蛛，火星”。结果展示在演示页面的底部，上方是生成的图像，下方是检测到的条件：

您还可以在面板中进一步细化您的配置：

Uni-ControlNet 也能很好地处理多条件设置。这里是一个结合了两种局部条件的例子：风暴兵的 Canny 边缘和森林的深度图。提示设置为“风暴兵在森林中的讲座”，结果如下：

借助 Uni-ControlNet，您甚至可以进一步结合更多条件。例如，您可以提供鹿、沙发、森林的局部条件，以及雪的全局条件，来创造一个不太可能自然发生的场景。提示设置为“森林中的沙发和鹿”，结果如下：

☕️ ：训练

首先，您需要下载 Stable Diffusion 的预训练权重 并将其放置在 ./ckpt/ 文件夹中。然后，您可以通过以下命令获取用于训练的初始权重：

python utils/prepare_weights.py init_local ckpt/v1-5-pruned.ckpt configs/local_v15.yaml ckpt/init_local.ckpt

python utils/prepare_weights.py init_global ckpt/v1-5-pruned.ckpt configs/global_v15.yaml ckpt/init_global.ckpt

这四个参数分别为模式、预训练的 SD 权重、模型配置文件以及初始权重的输出路径。

为了准备训练数据，请确保它们被放置在 ./data/ 文件夹中，并按照以下方式组织：

data/
├── anno.txt
├── images/
├── conditions/
    ├── condition-1/
    ├── condition-2/
    ...
...

具体来说，您可以使用 ./annotator/ 中的条件检测器来提取条件。然后，您需要将原始图像放入 ./data/images/ 文件夹，将提取的条件放入 ./data/conditions/condition-N/ 文件夹。而 ./data/anno.txt 是标注文件，每行代表一个训练样本，分为两部分：1) 文件 ID 和 2) 标注。请确保 ./data/anno.txt、./data/images/ 和 ./data/conditions/condition-N/ 目录中的文件 ID 保持一致。

现在，您只需运行以下命令即可使用自己的数据进行训练：

python src/train/train.py

请注意，局部适配器和全局适配器必须分别进行训练。此外，您还可以在 ./src/train/train.py 和 ./configs/ 中自定义训练配置。

完成单独训练后，您需要通过以下命令将两个适配器整合在一起：

python utils/prepare_weights.py integrate path1 path2 configs/uni_v15.yaml path3

其中，path1 和 path2 分别指带有局部和全局适配器的已训练 SD 权重，path3 则表示 Uni-ControlNet 的输出路径。

🎉 ：致谢：

本仓库基于 ControlNet 构建，衷心感谢他们的杰出工作！

📖 ：引用

@article{zhao2023uni,
  title={Uni-ControlNet: All-in-One Control to Text-to-Image Diffusion Models},
  author={Zhao, Shihao and Chen, Dongdong and Chen, Yen-Chun and Bao, Jianmin and Hao, Shaozhe and Yuan, Lu and Wong, Kwan-Yee~K.},
  journal={Advances in Neural Information Processing Systems},
  year={2023}
}

Uni-ControlNet 快速上手指南

Uni-ControlNet 是一个基于 Stable Diffusion 的多条件可控图像生成模型。它允许在一个模型中灵活组合多种局部控制（如边缘、姿态、深度）和全局控制（如内容风格），无需为每种条件单独训练模型。

环境准备

• 系统要求：Linux / Windows / macOS
• 硬件要求：推荐 NVIDIA GPU (显存 ≥ 8GB)
• 前置依赖：
  • Python 3.8+
  • Conda (推荐用于环境管理)
  • Git

安装步骤

1. 创建并激活虚拟环境

使用项目提供的配置文件创建 Conda 环境：

conda env create -f environment.yaml
conda activate unicontrol

2. 下载预训练模型

下载官方预训练权重（基于 Stable Diffusion v1.5），并将其放入 ./ckpt/ 目录。

• Google Drive: 下载链接
• Hugging Face: 下载链接

放置完成后，目录结构应包含 ./ckpt/uni.ckpt。

提示：如果下载速度较慢，可尝试使用国内镜像加速工具或代理。

基本使用

启动 Gradio 演示界面

运行以下命令即可加载预训练权重并启动本地 Web 界面：

python src/test/test.py

操作说明

启动后，浏览器将自动打开演示页面（通常位于 http://127.0.0.1:7860）。

1. 上传源图像：上传一张图片，系统会自动检测其素描/边缘信息。
2. 输入提示词：在文本框中输入生成提示词（例如："Robot spider, mars"）。
3. 选择控制条件：
   • 局部条件：支持 Canny 边缘、MLSD 线条、HED 边界、素描、OpenPose 姿态、Midas 深度图、分割掩码等 7 种。
   • 全局条件：支持内容风格控制。
4. 多条件组合：你可以同时启用多个条件。例如，结合“风暴兵的边缘图”和“森林的深度图”，配合提示词 "Stormtrooper's lecture in the forest" 进行生成。
5. 生成结果：点击生成按钮，页面下方将展示生成的图像及对应的条件检测图。

通过面板配置，你可以进一步调整采样步数、引导尺度等参数以获得更佳效果。