UniControl

这个代码库对应的论文是：

UniControl: A Unified Diffusion Model for Controllable Visual Generation In the Wild
Can Qin ^1,2, Shu Zhang¹, Ning Yu ¹, Yihao Feng¹, Xinyi Yang¹, Yingbo Zhou ¹, Huan Wang ¹, Juan Carlos Niebles¹, Caiming Xiong ¹, Silvio Savarese ¹, Stefano Ermon ³, Yun Fu ², Ran Xu ¹
¹ Salesforce AI ² Northeastern University ³ Stanford University
当Can Qin在Salesforce AI Research实习时完成的工作。

介绍

我们介绍了UniControl，一个新的生成基础模型，它在一个框架内整合了广泛的可控条件到图像（C2I）任务，同时仍然允许任意的语言提示。UniControl实现了像素级精确的图像生成，视觉条件主要影响生成的结构，语言提示引导风格和背景。为了让UniControl具备处理多样化视觉条件的能力，我们增强了预训练的文本-图像扩散模型，并引入了一个任务感知的HyperNet来调节扩散模型，使其能够同时适应不同的C2I任务。实验结果表明，UniControl通常表现优于具有相似模型大小的单任务控制方法。这种控制的多样性使UniControl在可控视觉生成领域中成为一项重要的进步。

更新

05/18/23: UniControl 论文上传到 arXiv.
05/26/23: UniControl 推理代码和检查点公开。
05/28/23: 最新的UniControl模型检查点 (1.4B 参数, 5.78GB) 更新。
06/08/23: 最新的UniControl模型检查点更新，现在支持12个任务（Canny, HED, Sketch, Depth, Normal, Skeleton, Bbox, Seg, Outpainting, Inpainting, Deblurring 和 Colorization) !
06/08/23: 训练数据集 (MultiGen-20M) 完全发布。
06/08/23: 训练代码公开。:blush:
07/06/23: 最新的UniControl模型 v1.1 检查点更新，现在支持12个任务（Canny, HED, Sketch, Depth, Normal, Skeleton, Bbox, Seg, Outpainting, Inpainting, Deblurring 和 Colorization) !
07/25/23: Huggingface Demo API 可用了！
07/25/23: Safetensors 模型可用了！检查点
09/21/23: 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 Demo (应用示例视频, 需要CUDA 11.0和Conda 4.12.0)

我们为不同任务提供了Gradio demos。示例图像保存在 ./test_imgs。

对于所有任务 (Canny, HED, Sketch, Depth, Normal, Human Pose, Seg, Bbox, Outpainting, Colorization, Deblurring, Inpainting) 请运行以下代码：

python app/gradio_all_tasks.py

我们支持直接从条件到图像的生成（如上图所示）。如果您希望直接上传条件图像，请在UI中取消选中条件提取。

或者，我们提供特定任务的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

待办事项

提示

负面提示有时非常有用: 黑白,低分辨率,糟糕的解剖结构,最差的质量,低质量 是负面提示的例子。
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