使用 🧨 Diffusers 实现超快速 ControlNet

发布于 2023 年 3 月 3 日
在 GitHub 上更新
Open In Colab

自 Stable Diffusion 风靡全球以来,人们一直在寻找方法来更好地控制生成过程的结果。ControlNet 提供了一个最小的接口,允许用户在很大程度上自定义生成过程。通过 ControlNet,用户可以轻松地通过不同的空间上下文(如深度图、分割图、涂鸦、关键点等)来控制生成!

我们可以将卡通画变成具有惊人连贯性的真实照片。

逼真的 Lofi 女孩

甚至可以用它来做你的室内设计师。

之前 之后

您可以将您的草图涂鸦变成艺术画作。

之前 之后

此外,让一些著名标志栩栩如生。

之前 之后

有了 ControlNet,一切皆有可能 🌠

在这篇博文中,我们首先介绍 StableDiffusionControlNetPipeline,然后展示它如何应用于各种控制条件。让我们开始控制吧!

ControlNet: 简而言之

ControlNet 由 Lvmin Zhang 和 Maneesh Agrawala 在 Adding Conditional Control to Text-to-Image Diffusion Models 中引入。它引入了一个框架,允许支持各种空间上下文,这些上下文可以作为 Diffusion 模型(如 Stable Diffusion)的额外条件。Diffusers 的实现改编自原始 源代码

训练 ControlNet 包含以下步骤:

  1. 克隆 Diffusion 模型的预训练参数(例如 Stable Diffusion 的潜在 UNet),(称为“可训练副本”)同时单独维护预训练参数(“锁定副本”)。这样做是为了让锁定参数副本能够保留从大量数据集中学习到的丰富知识,而可训练副本则用于学习特定任务的方面。
  2. 参数的可训练副本和锁定副本通过“零卷积”层连接(更多信息请参阅此处),这些层作为 ControlNet 框架的一部分进行优化。这是一种训练技巧,用于在训练新条件时保留冻结模型已学习的语义。

从图像上看,训练 ControlNet 如下所示:

controlnet-structure
该图取自此处

ControlNet 训练的训练集样本如下所示(通过边缘图进行额外条件控制)

提示 原始图片 条件
"鸟"

同样,如果我们要用语义分割图来控制 ControlNet,一个训练样本会是这样:

提示 原始图片 条件
"大房子"

每种新的条件类型都需要训练一个新的 ControlNet 权重副本。论文提出了 8 种不同的条件模型,这些模型在 Diffusers 中都受支持

对于推理,需要预训练的扩散模型权重和训练好的 ControlNet 权重。例如,使用 Stable Diffusion v1-5 和 ControlNet 检查点比仅使用原始 Stable Diffusion 模型需要多出大约 7 亿个参数,这使得 ControlNet 在推理时对内存的消耗更大。

由于预训练的扩散模型在训练期间被锁定,因此在使用不同条件时只需要切换 ControlNet 参数。这使得在单个应用程序中部署多个 ControlNet 权重变得相当简单,如下所示。

StableDiffusionControlNetPipeline

在开始之前,我们要向社区贡献者 Takuma Mori 致以崇高的敬意,感谢他领导了 ControlNet 在 Diffusers 中的集成 ❤️。

为了体验 ControlNet,Diffusers 提供了 StableDiffusionControlNetPipeline,类似于其他 Diffusers 管线StableDiffusionControlNetPipeline 的核心是 controlnet 参数,它允许我们在保持预训练扩散模型权重不变的情况下提供一个特定的已训练 ControlNetModel 实例。

我们将在本博客文章中探索 StableDiffusionControlNetPipeline 的不同用例。我们将介绍的第一个 ControlNet 模型是 Canny 模型——这是最受欢迎的模型之一,它生成了一些您可能在互联网上看到的令人惊叹的图像。

欢迎您使用 此 Colab 笔记本 运行以下部分中显示的代码片段。

在开始之前,我们先确保安装了所有必要的库

pip install diffusers==0.14.0 transformers xformers git+https://github.com/huggingface/accelerate.git

为了根据选择的 ControlNet 处理不同的条件,我们还需要安装一些额外的依赖项

pip install opencv-contrib-python
pip install controlnet_aux

我们将使用著名的画作 《戴珍珠耳环的少女》 作为这个例子。所以,让我们下载图片并查看一下

from diffusers.utils import load_image

image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
)
image

接下来,我们将把图像通过 canny 预处理器处理

import cv2
from PIL import Image
import numpy as np

image = np.array(image)

low_threshold = 100
high_threshold = 200

image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
canny_image

如我们所见,它本质上是边缘检测

现在,我们加载 runwaylml/stable-diffusion-v1-5 以及用于 Canny 边缘的 ControlNet 模型。模型以半精度 (torch.dtype) 加载,以实现快速且内存高效的推理。

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch

controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)

我们不使用 Stable Diffusion 默认的 PNDMScheduler,而是使用目前最快的扩散模型调度器之一,称为 UniPCMultistepScheduler。选择改进的调度器可以大幅减少推理时间——在我们的例子中,我们可以将推理步数从 50 步减少到 20 步,同时或多或少地保持相同的图像生成质量。有关调度器的更多信息可以在此处找到。

from diffusers import UniPCMultistepScheduler

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

我们没有直接将管道加载到 GPU,而是启用了智能 CPU 卸载,这可以通过 enable_model_cpu_offload 函数实现。

请记住,在推理扩散模型(如 Stable Diffusion)时,不仅需要一个,而且需要多个模型组件按顺序运行。对于带有 ControlNet 的 Stable Diffusion,我们首先使用 CLIP 文本编码器,然后是扩散模型 UNet 和 ControlNet,然后是 VAE 解码器,最后运行安全检查器。大多数组件在扩散过程中只运行一次,因此不需要一直占用 GPU 内存。通过启用智能模型卸载,我们确保每个组件只在需要时才加载到 GPU 中,这样我们可以在不显著减慢推理速度的情况下显著节省内存消耗。

注意:运行 enable_model_cpu_offload 时,请勿手动使用 .to("cuda") 将管道移动到 GPU——一旦启用 CPU 卸载,管道会自动处理 GPU 内存管理。

pipe.enable_model_cpu_offload()

最后,我们要充分利用令人惊叹的 FlashAttention/xformers 注意力层加速功能,所以让我们启用它!如果此命令对您不起作用,则可能没有正确安装 xformers。在这种情况下,您可以跳过以下代码行。

pipe.enable_xformers_memory_efficient_attention()

现在我们准备运行 ControlNet 管道了!

我们仍然提供一个提示来指导图像生成过程,就像我们通常使用 Stable Diffusion 图像到图像管道那样。然而,ControlNet 将对生成的图像提供更多的控制,因为我们将能够使用我们刚刚创建的 canny 边缘图像来控制生成图像中的精确构图。

看看一些当代名人摆出这幅 17 世纪画作的相同姿势的图片会很有趣。使用 ControlNet 做到这一点非常容易,我们所要做的就是将这些名人的名字包含在提示中!

首先,我们创建一个简单的辅助函数来以网格形式显示图像。

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

接下来,我们定义输入提示并设置一个种子以确保可重现性。

prompt = ", best quality, extremely detailed"
prompt = [t + prompt for t in ["Sandra Oh", "Kim Kardashian", "rihanna", "taylor swift"]]
generator = [torch.Generator(device="cpu").manual_seed(2) for i in range(len(prompt))]

最后,我们可以运行管道并显示图像!

output = pipe(
    prompt,
    canny_image,
    negative_prompt=["monochrome, lowres, bad anatomy, worst quality, low quality"] * 4,
    num_inference_steps=20,
    generator=generator,
)

image_grid(output.images, 2, 2)

我们还可以轻松地将 ControlNet 与微调结合起来!例如,我们可以使用 DreamBooth 微调模型,并用它将我们自己渲染到不同的场景中。

在这篇文章中,我们将以我们心爱的土豆先生为例,展示如何将 ControlNet 与 DreamBooth 结合使用。

我们可以使用相同的 ControlNet。但是,我们不使用 Stable Diffusion 1.5,而是将 土豆先生模型 加载到我们的管道中——土豆先生是一个使用 Dreambooth 🥔 微调的 Stable Diffusion 模型,它使用土豆先生的概念。

让我们再次运行上面的命令,保持相同的 ControlNet!

model_id = "sd-dreambooth-library/mr-potato-head"
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()

现在,让我们让土豆先生为 约翰内斯·维米尔 摆姿势吧!

generator = torch.manual_seed(2)
prompt = "a photo of sks mr potato head, best quality, extremely detailed"
output = pipe(
    prompt,
    canny_image,
    negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
    num_inference_steps=20,
    generator=generator,
)
output.images[0]

可以看出,土豆先生不是最好的候选人,但他尽力了,在捕捉一些精髓方面做得相当不错 🍟

ControlNet 的另一个独家应用是,我们可以从一张图像中提取姿势,并将其重复使用以生成具有完全相同姿势的不同图像。因此,在下一个示例中,我们将使用 Open Pose ControlNet 教超级英雄们如何做瑜伽!

首先,我们需要获取一些正在做瑜伽的人的图像

urls = "yoga1.jpeg", "yoga2.jpeg", "yoga3.jpeg", "yoga4.jpeg"
imgs = [
    load_image("https://huggingface.co/datasets/YiYiXu/controlnet-testing/resolve/main/" + url) 
    for url in urls
]

image_grid(imgs, 2, 2)

现在,让我们使用 controlnet_aux 中方便提供的 OpenPose 预处理器提取瑜伽姿势。

from controlnet_aux import OpenposeDetector

model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")

poses = [model(img) for img in imgs]
image_grid(poses, 2, 2)

为了使用这些瑜伽姿势生成新图像,我们创建一个 Open Pose ControlNet。我们将生成一些超级英雄图像,但采用上面所示的瑜伽姿势。让我们开始吧 🚀

controlnet = ControlNetModel.from_pretrained(
    "fusing/stable-diffusion-v1-5-controlnet-openpose", torch_dtype=torch.float16
)

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

现在是瑜伽时间!

generator = [torch.Generator(device="cpu").manual_seed(2) for i in range(4)]
prompt = "super-hero character, best quality, extremely detailed"
output = pipe(
    [prompt] * 4,
    poses,
    negative_prompt=["monochrome, lowres, bad anatomy, worst quality, low quality"] * 4,
    generator=generator,
    num_inference_steps=20,
)
image_grid(output.images, 2, 2)

结合多种条件

单个图像生成可以结合多个 ControlNet 条件。将 ControlNet 列表传递给管道的构造函数,并将相应的条件列表传递给 __call__

在结合条件时,遮盖条件使其不重叠是很有帮助的。在这个例子中,我们遮盖了 canny 映射的中间部分,其中包含姿势条件。

改变 controlnet_conditioning_scale 也可以帮助强调一个条件而不是另一个。

Canny 条件

原始图片

准备条件

from diffusers.utils import load_image
from PIL import Image
import cv2
import numpy as np
from diffusers.utils import load_image

canny_image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/landscape.png"
)
canny_image = np.array(canny_image)

low_threshold = 100
high_threshold = 200

canny_image = cv2.Canny(canny_image, low_threshold, high_threshold)

# zero out middle columns of image where pose will be overlayed
zero_start = canny_image.shape[1] // 4
zero_end = zero_start + canny_image.shape[1] // 2
canny_image[:, zero_start:zero_end] = 0

canny_image = canny_image[:, :, None]
canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
canny_image = Image.fromarray(canny_image)

Openpose 条件

原始图片

准备条件

from controlnet_aux import OpenposeDetector
from diffusers.utils import load_image

openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")

openpose_image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/person.png"
)
openpose_image = openpose(openpose_image)

运行带有多个条件的 ControlNet

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import torch

controlnet = [
    ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16),
    ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16),
]

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

pipe.enable_xformers_memory_efficient_attention()
pipe.enable_model_cpu_offload()

prompt = "a giant standing in a fantasy landscape, best quality"
negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"

generator = torch.Generator(device="cpu").manual_seed(1)

images = [openpose_image, canny_image]

image = pipe(
    prompt,
    images,
    num_inference_steps=20,
    generator=generator,
    negative_prompt=negative_prompt,
    controlnet_conditioning_scale=[1.0, 0.8],
).images[0]

image.save("./multi_controlnet_output.png")

在这些示例中,我们探索了 StableDiffusionControlNetPipeline 的多个方面,以展示通过 Diffusers 使用 ControlNet 是多么容易和直观。但是,我们没有涵盖 ControlNet 支持的所有条件类型。要了解更多信息,我们鼓励您查看相应的模型文档页面:

我们欢迎您组合这些不同的元素,并与 @diffuserslib 分享您的成果。务必查看Colab 笔记本,体验上述一些示例!

我们还展示了一些技术,通过使用快速调度器、智能模型卸载和 xformers 来加快生成过程并节省内存。结合这些技术,生成过程在 V100 GPU 上仅需约 3 秒,单张图像仅消耗约 4 GB VRAM ⚡️ 在 Google Colab 等免费服务上,生成时间约为 5 秒(默认 GPU 为 T4),而原始实现需要 17 秒才能创建相同结果!将 diffusers 工具箱中的所有组件结合起来真是一项超能力 💪

总结

我们一直在大量使用 StableDiffusionControlNetPipeline,到目前为止我们的体验非常有趣!我们很高兴看到社区在这个管道之上构建什么。如果您想查看 Diffusers 中支持的其他允许受控生成的管道和技术,请查看我们的 官方文档

如果您迫不及待地想直接尝试 ControlNet,我们也为您提供了!只需单击以下空间之一即可玩转 ControlNet

  • Canny ControlNet Spaces
  • OpenPose ControlNet Spaces

社区

注册登录 发表评论