使用 ? Diffusers 實現 ControlNet 高速推理

HuggingFace發表於2023-03-07

自從 Stable Diffusion 風靡全球以來,人們一直在尋求如何更好地控制生成過程的方法。ControlNet 提供了一個簡單的遷移學習方法,能夠允許使用者在很大程度上自定義生成過程。透過 ControlNet,使用者可以輕鬆地使用多種空間語義條件資訊 (例如深度圖、分割圖、塗鴉圖、關鍵點等) 來控制生成過程。

具體來說,我們可以:

將卡通繪圖轉化為逼真的照片,同時保持極佳的佈局連貫性。

寫實版的 Lofi Girl (上: 原圖,下: 新圖)

寫實版的 Lofi Girl (上: 原圖,下: 新圖)

進行室內設計。

原圖新圖

將塗鴉草圖變成藝術作品。

原圖新圖

甚至擬人化著名的 logo 形象。

原圖新圖

ControlNet,使一切皆有可能 ?

本文的主要內容:

  1. 介紹 StableDiffusionControlNetPipeline
  2. 展示多種控制條件樣例

讓我們開啟控制之旅!

ControlNet 簡述

ControlNet 在 Adding Conditional Control to Text-to-Image Diffusion Models 一文中提被出,作者是 Lvmin Zhang 和 Maneesh Agrawala。它引入了一個框架,支援在擴散模型 (如 Stable Diffusion) 上附加額外的多種空間語義條件來控制生成過程。

訓練 ControlNet 包括以下步驟:

  1. 克隆擴散模型的預訓練引數 (文中稱為 可訓練副本, trainable copy。如 Stable Diffusion 的 latent UNet 部分),同時保留原本的預訓練引數 (文中稱為 鎖定副本, locked copy)。這樣可以實現: a) 讓鎖定副本保留從大型資料集中學到的豐富知識;b) 讓可訓練副本學習特定任務的知識。
  2. 可訓練副本和鎖定副本的引數透過 “零卷積” 層 (詳見 此處) 連線。“零卷積” 層是 ControlNet 框架的一部分,會在特定任務中最佳化引數。這是一種訓練技巧,可以在新任務條件訓練時保留已凍結模型已經學到的語義資訊。

訓練 ControlNet 的過程如圖所示:

ControlNet 訓練集中的其中一種樣例如下 (額外的控制條件是 Canny 邊緣圖):

原始圖片應用控制條件

同樣地,如果我們使用的額外控制條件是語義分割圖,那麼 ControlNet 訓練集的樣例就是這樣:

原始圖片應用控制條件

每對 ControlNet 施加一種額外的控制條件,都需要訓練一份新的可訓練副本引數。論文中提出了 8 種不同的控制條件,對應的控制模型在 Diffusers 中均已支援!

推理階段需要同時使用擴散模型的預訓練權重以及訓練過的 ControlNet 權重。如要使用 Stable Diffusion v1-5 以及其 ControlNet 權重推理,其引數量要比僅使用 Stable Diffusion v1-5 多大約 7 億個,因此推理 ControlNet 需要消耗更多的記憶體。

由於在訓練過程中擴散模型預訓練引數為鎖定副本,因此在使用不同的控制條件訓練時,只需要切換 ControlNet 可訓練副本的引數即可。這樣在一個應用程式中部署多個 ControlNet 權重就非常簡單了,本文會在後面詳細介紹。

StableDiffusionControlNetPipeline

在開始之前,我們要向社群貢獻者 Takuma Mori (@takuma104) 表示巨大的感謝。將 ControlNet 整合到 Diffusers 中,他功不可沒 ❤️。

類似 Diffusers 中的 其他 Pipeline,Diffusers 同樣為 ControlNet 提供了 StableDiffusionControlNetPipeline 供使用者使用。StableDiffusionControlNetPipeline 的核心是 controlnet 引數,它接收使用者指定的訓練過的 ControlNetModel 例項作為輸入,同時保持擴散模型的預訓練權重不變。

本文將介紹 StableDiffusionControlNetPipeline 的多個不同用例。首先要介紹的第一個 ControlNet 模型是 Canny 模型,這是目前最流行的 ControlNet 模型之一,您可能已經在網上見識過一些它生成的精美圖片。在閱讀到各個部分的程式碼時,也歡迎您使用此 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://hf.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

如圖可見,Canny 本質上是邊緣檢測器:

接下來,我們載入 runwaylml/stable-diffusion-v1-5Canny 邊緣 ControlNet 模型。設定引數 torch.dtype=torch.float16 可以指定模型以半精度模式載入,可實現記憶體高效和快速的推理。

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)

我們透過呼叫 enable_model_cpu_offload 函式來啟用智慧 CPU 解除安裝,而不是直接將 pipeline 載入到 GPU 上。

智慧 CPU 解除安裝是一種降低視訊記憶體佔用的方法。擴散模型 (如 Stable Diffusion) 的推理並不是執行一個單獨的模型,而是多個模型元件的序列推理。如在推理 ControlNet Stable Diffusion 時,需要首先執行 CLIP 文字編碼器,其次推理擴散模型 UNet 和 ControlNet,然後執行 VAE 解碼器,最後執行 safety checker (安全檢查器,主要用於稽核過濾違規影像)。而在擴散過程中大多陣列件僅執行一次,因此不需要一直佔用 GPU 記憶體。透過啟用智慧模型解除安裝,可以確保每個元件在不需要參與 GPU 計算時解除安裝到 CPU 上,從而顯著降低視訊記憶體佔用,並且不會顯著增加推理時間 (僅增加了模型在 GPU-CPU 之間的轉移時間)。

注意: 啟用 enable_model_cpu_offload 後,pipeline 會自動進行 GPU 記憶體管理,因此請不要再使用 .to("cuda") 手動將 pipeline 轉移到 GPU。

pipe.enable_model_cpu_offload()

最後,我們要充分利用 FlashAttention/xformers 進行注意力層加速。執行下列程式碼以實現加速,如果該程式碼沒有起作用,那麼您可能沒有正確安裝 xformers 庫,此時您可以跳過該程式碼。

pipe.enable_xformers_memory_efficient_attention()

基本條件準備就緒,現在來執行 ControlNet pipeline!

跟執行 Stable Diffusion image-to-image pipeline 相同的是,我們也使用了文字提示語來引導影像生成過程。不過有一些不同的是,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))]

最後執行 pipeline,並視覺化生成的影像!

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 增加控制資訊,將其渲染到不同的場景中。

本文將以我們最愛的土豆先生為例,來介紹怎樣結合使用 ControlNet 和 DreamBooth。

相較於上文,pipeline 中使用的 ControlNet 部分保持不變,但是不使用 Stable Diffusion 1.5,而是重新載入一個 土豆先生 模型 (使用 Dreambooth 微調的 Stable Diffusion 模型) ?。

雖然 ControlNet 沒變,但仍然需要重新載入 pipeline。

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 pipeline 來生成一些相同姿態的超級英雄影像。Let's go ?

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)

透過以上示例,我們對 StableDiffusionControlNetPipeline 的多種用法有了直觀的認識,也學會了如何使用 Diffusers 玩轉 ControlNet。不過,還有一些 ControlNet 支援的其他型別的控制條件示例,由於篇幅原因本文不再展開,如想了解更多資訊,可以點選以下連結檢視相應的模型文件頁面:

我們非常歡迎您嘗試組合不同的控制元件來生成精美的影像,並在 Twitter 上與 @diffuserslib 分享您的作品。如果您還沒有執行上述程式碼段,這裡再次建議您檢視剛才提到的 Colab 筆記本,親自執行程式碼體驗示例的效果!

在上文中,我們介紹了加速生成過程、減少視訊記憶體佔用的一些技巧,它們包括: 快速排程器、智慧模型解除安裝、xformers。如果結合使用這些技巧,單張影像的生成過程僅需要: V100 GPU 上約 3 秒的推理時間以及約 4 GB 的 VRAM 佔用;免費 GPU 服務 (如 Google Colab 的 T4) 上約 5 秒的推理時間。如果沒有實現這些技巧,同樣的生成過程可達 17 秒!現已整合至 Diffusers 工具箱,來使用 Diffusers 吧,它真的非常強力!?

結語

本文介紹了 StableDiffusionControlNetPipeline 的多個用例,非常有趣!我們也非常期待看到社群在此 pipeline 的基礎上能構建出什麼好玩的應用。如果您想了解更多 Diffusers 支援的關於控制模型的其他 pipeline 和技術細節,請檢視我們的 官方文件

如果您想直接嘗試 ControlNet 的控制效果,我們也能滿足!只需點選以下 HuggingFace Spaces 即可嘗試控制生成影像:


原文連結: https://huggingface.co/blog/controlnet

作者: Sayak Paul、YiYi Xu、Patrick von Platen

譯者: SuSung-boy

審校、排版: zhongdongy (阿東)

相關文章