OpenVINO賦能潛在一致性模型(LCMs)的高速圖像生成

作者：武卓英特爾 AI 布道師

提起生成式 AI，從去年開(kāi)始爆火至今，相信無(wú)論是不是 AI 領(lǐng)域的開(kāi)發(fā)者，對(duì)這個(gè)概念都是毫不陌生了。說(shuō)到生成式 AI 領(lǐng)域的典型應(yīng)用場(chǎng)景和模型，相信大家最先想到的就是文生圖這個(gè)場(chǎng)景，以及它背后一系列的潛在擴(kuò)散模型（Latent Diffusion models, LDMs）。借由這些模型的強(qiáng)大能力，只需要輸入一段文字，人人都可化身設(shè)計(jì)師、插畫(huà)師，進(jìn)行精美圖片的創(chuàng)作。但是，調(diào)用文生圖模型的時(shí)候，若是選擇調(diào)用云端運(yùn)行的 API ，可能需要付費(fèi)的同時(shí)、還有可能面臨排隊(duì)等待等問(wèn)題。如果選擇在本地機(jī)器上運(yùn)行，那么對(duì)于機(jī)器的算力和內(nèi)存就有較高的要求，同時(shí)也需要進(jìn)行長(zhǎng)時(shí)間的等待，畢竟這些模型往往需要迭代幾十次之后才能生成一幅比較精美的圖片。

最近一個(gè)稱(chēng)之為LCMs（Latent Consistency Models）的模型橫空出世，讓文生圖模型的圖片急速生成成為了可能。受到一致性模型（CM）的啟發(fā)，潛在一致性模型（LCMs），可以在任何預(yù)訓(xùn)練的潛在擴(kuò)散模型上進(jìn)行快速推斷、步驟最少，包括穩(wěn)定擴(kuò)散模型（Stable Diffusion）。一致性模型是一種新的生成模型家族，可以實(shí)現(xiàn)一步或少量步驟生成。其核心思想是學(xué)習(xí) PF-ODE（常微分方程概率流的軌跡）解的函數(shù)。通過(guò)學(xué)習(xí)保持 ODE 軌跡上點(diǎn)的一致性的一致性映射，這些模型允許進(jìn)行單步生成，消除了計(jì)算密集的迭代的需求。然而，CM 受限于像素空間圖像生成任務(wù)，因此不適用于合成高分辨率圖像。LCMs 在圖像潛在空間采用一致性模型來(lái)生成高分辨率圖像。將引導(dǎo)反向擴(kuò)散過(guò)程視為解決 PF-ODE 的過(guò)程。LCMs 旨在直接預(yù)測(cè)潛在空間中這種 ODE 的解，減少了大量迭代的需求，實(shí)現(xiàn)了快速、高保真度的采樣。在大規(guī)模擴(kuò)散模型（如 Stable Diffusion）中利用圖像潛在空間有效地增強(qiáng)了圖像生成質(zhì)量并減少了計(jì)算負(fù)載，使得圖片急速生成成為了可能。

有關(guān)所提出方法和模型的更多詳細(xì)信息，可以在項(xiàng)目頁(yè)面[1]、論文[2]和原始倉(cāng)庫(kù)[3]中找到。

如此充滿(mǎn)魔力的 LCMs 文生圖模型，我們的 OpenVINO 當(dāng)然也可以對(duì)它進(jìn)行完全的優(yōu)化、壓縮以及推理加速、快速部署的支持。接下來(lái)，就讓我們通過(guò)我們常用的 OpenVINO Notebooks 倉(cāng)庫(kù)中關(guān)于 LCMs 模型的 Jupyter Notebook 代碼[4] 和拆解，來(lái)進(jìn)一步了解具體步驟吧。

第一步：

安裝相應(yīng)工具包、加載模型

并轉(zhuǎn)換為 OpenVINO IR 格式

%pip install -q "torch" --index-url https://download.pytorch.org/whl/cpu
%pip install -q "openvino>=2023.1.0" transformers "diffusers>=0.22.0" pillow gradio "nncf>=2.6.0" datasets

左滑查看更多

模型下載

與傳統(tǒng)的 Stable Diffusion 流水線類(lèi)似， LCMs 模型內(nèi)部也包含了 Text Encoder, U-Net, VAE Decoder 三個(gè)模型。

Text Encoder 模型

負(fù)責(zé)從文本 text prompt創(chuàng)建圖片生成的條件

U-Net 模型

負(fù)責(zé)初步降噪潛在圖像表示

Autoencoder (VAE) Decoder

用于將潛在空間解碼為最終的圖片

因此分別需要對(duì)這三個(gè)模型進(jìn)行下載，部分代碼如下：

import gc
import warnings
from pathlib import Path
from diffusers import DiffusionPipeline


 
warnings.filterwarnings("ignore")


TEXT_ENCODER_OV_PATH = Path("model/text_encoder.xml")
UNET_OV_PATH = Path("model/unet.xml")
VAE_DECODER_OV_PATH = Path("model/vae_decoder.xml")
 


 def load_orginal_pytorch_pipeline_componets(skip_models=False, skip_safety_checker=True):

左滑查看更多

skip_conversion = (
TEXT_ENCODER_OV_PATH.exists()
and UNET_OV_PATH.exists()
and VAE_DECODER_OV_PATH.exists()
)


(
scheduler,
tokenizer,
feature_extractor,
safety_checker,
text_encoder,
unet,
vae,
) = load_orginal_pytorch_pipeline_componets(skip_conversion)

左滑查看更多

模型轉(zhuǎn)換

包括將以上三個(gè)模型的轉(zhuǎn)換為 OpenVINO IR 格式。

第二步：

準(zhǔn)備基于 OpenVINO

的推理流水線

流水線如下圖所示。

整個(gè)流水線利用一個(gè)潛在圖像表示和一個(gè)文本提示，通過(guò) CLIP 的文本編碼器將文本提示轉(zhuǎn)化為文本嵌入作為輸入。初始的潛在圖像表示是使用隨機(jī)噪聲生成器生成的。與原始的 Stable Diffusion 流程不同，LCMs 還使用引導(dǎo)尺度來(lái)獲取時(shí)間步驟條件嵌入作為擴(kuò)散過(guò)程的輸入，而在 Stable Diffusion 中，它用于縮放輸出的潛在表示。

接下來(lái)，U-Net 迭代地去噪隨機(jī)潛在圖像表示，同時(shí)以文本嵌入為條件。U-Net 的輸出是噪聲殘差，通過(guò)調(diào)度算法用于計(jì)算去噪后的潛在圖像表示。LCMs 引入了自己的調(diào)度算法，擴(kuò)展了去噪擴(kuò)散概率模型（DDPMs）中引入的非馬爾可夫引導(dǎo)。去噪過(guò)程會(huì)重復(fù)多次（注意：原始 SD 流程中默認(rèn)為 50 次，但對(duì)于 LCM，只需要 2-8 次左右的小步驟即可！），以逐步獲取更好的潛在圖像表示。完成后，潛在圖像表示將由變 VAE 解碼器解碼為最終的圖像輸出。

定義關(guān)于 LCMs 推理流水線的類(lèi)，部分代碼如下：

from typing import Union, Optional, Any, List, Dict
from transformers import CLIPTokenizer, CLIPImageProcessor
from diffusers.pipelines.stable_diffusion.safety_checker import (
  StableDiffusionSafetyChecker,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.image_processor import VaeImageProcessor


 
class OVLatentConsistencyModelPipeline(DiffusionPipeline):

左滑查看更多

第三步：

配置推理流水線

首先，創(chuàng)建 OpenVINO模型的實(shí)例，并使用選定的設(shè)備對(duì)其進(jìn)行編譯。從下拉列表中選擇設(shè)備，以便使用 OpenVINO運(yùn)行推理。

core = ov.Core()


import ipywidgets as widgets


device = widgets.Dropdown(
  options=core.available_devices + ["AUTO"],
  value="CPU",
  description="Device:",
  disabled=False,
)


device

左滑查看更多

text_enc = core.compile_model(TEXT_ENCODER_OV_PATH, device.value)
unet_model = core.compile_model(UNET_OV_PATH, device.value)


ov_config = {"INFERENCE_PRECISION_HINT": "f32"} if device.value != "CPU" else {}


vae_decoder = core.compile_model(VAE_DECODER_OV_PATH, device.value, ov_config)

左滑查看更多

模型分詞器（tokenizer）和調(diào)度器（scheduler）也是流水線的重要組成部分。該流水線還可以使用安全檢查器，該過(guò)濾器用于檢測(cè)相應(yīng)生成的圖像是否包含“不安全工作”（nsfw）內(nèi)容。nsfw 內(nèi)容檢測(cè)過(guò)程需要使用 CLIP 模型獲得圖像嵌入，因此需要在流水線中添加額外的特征提取器組件。我們重用原始 LCMs 流水線中的標(biāo)記器、特征提取器、調(diào)度器和安全檢查器。

ov_pipe = OVLatentConsistencyModelPipeline(
  tokenizer=tokenizer,
  text_encoder=text_enc,
  unet=unet_model,
  vae_decoder=vae_decoder,
  scheduler=scheduler,
  feature_extractor=feature_extractor,
  safety_checker=safety_checker,
)

左滑查看更多

第四步：

文本到圖片生成

prompt = "a beautiful pink unicorn, 8k"
num_inference_steps = 4
torch.manual_seed(1234567)


images = ov_pipe(
  prompt=prompt,
  num_inference_steps=num_inference_steps,
  guidance_scale=8.0,
  lcm_origin_steps=50,
  output_type="pil",
  height=512,
  width=512,
).images

左滑查看更多

在我本地的機(jī)器上，我分別使用了我的英特爾 酷睿 第 12 代 CPU、以及英特爾銳炫 顯卡A770m運(yùn)行了模型推理，生成圖片真的是在眨眼之間！

當(dāng)然，為了方便各位開(kāi)發(fā)者的使用，我們的notebook代碼[5] 中還為大家設(shè)計(jì)了基于 Gradio 的、更加用戶(hù)優(yōu)化的界面。

第五步：

利用 NNCF 對(duì)模型進(jìn)行量化壓縮

另外，如果大家對(duì)于模型還有進(jìn)一步的尺寸壓縮以及內(nèi)存占用壓縮的需求，我們的 notebook 中也同樣提供了基于 NNCF 進(jìn)行量化壓縮的代碼示例。

量化壓縮的過(guò)程分為以下三個(gè)步驟：

為量化創(chuàng)建一個(gè)校準(zhǔn)數(shù)據(jù)集

運(yùn)行nncf.quantize()獲得量化模型

利用 openvino.save_model() 保存量化為INT8格式的模型

部分代碼如下：

%%skip not $to_quantize.value
import nncf
from nncf.scopes import IgnoredScope


if UNET_INT8_OV_PATH.exists():
  print("Loading quantized model")
  quantized_unet = core.read_model(UNET_INT8_OV_PATH)
else:
  unet = core.read_model(UNET_OV_PATH)
  quantized_unet = nncf.quantize(
    model=unet,
    subset_size=subset_size,
    preset=nncf.QuantizationPreset.MIXED,
    calibration_dataset=nncf.Dataset(unet_calibration_data),
    model_type=nncf.ModelType.TRANSFORMER,
    advanced_parameters=nncf.AdvancedQuantizationParameters(
      disable_bias_correction=True
    )
  )
  ov.save_model(quantized_unet, UNET_INT8_OV_PATH)

左滑查看更多

運(yùn)行效果如下：

根據(jù)同樣的文本 text prompt，檢測(cè)一下量化后模型的生成效果：

當(dāng)然推理時(shí)間上的表現(xiàn)有何提升呢，我們也歡迎大家利用如下的代碼在自己的機(jī)器上實(shí)測(cè)一下：

%%skip not $to_quantize.value


import time


validation_size = 10
calibration_dataset = datasets.load_dataset("laion/laion2B-en", split="train", streaming=True).take(validation_size)
validation_data = []
while len(validation_data) < validation_size:
 ? ?batch = next(iter(calibration_dataset))
 ? ?prompt = batch["TEXT"]
 ? ?validation_data.append(prompt)


def calculate_inference_time(pipeline, calibration_dataset):
 ? ?inference_time = []
 ? ?pipeline.set_progress_bar_config(disable=True)
 ? ?for prompt in calibration_dataset:
 ? ? ? ?start = time.perf_counter()
 ? ? ? ?_ = pipeline(
 ? ? ? ? ? ?prompt,
 ? ? ? ? ? ?num_inference_steps=num_inference_steps,
 ? ? ? ? ? ?guidance_scale=8.0,
 ? ? ? ? ? ?lcm_origin_steps=50,
 ? ? ? ? ? ?output_type="pil",
 ? ? ? ? ? ?height=512,
 ? ? ? ? ? ?width=512,
 ? ? ? ?)
 ? ? ? ?end = time.perf_counter()
 ? ? ? ?delta = end - start
 ? ? ? ?inference_time.append(delta)
 ? ?return np.median(inference_time)

左滑查看更多

小結(jié)

整個(gè)的步驟就是這樣！現(xiàn)在就開(kāi)始跟著我們提供的代碼和步驟，動(dòng)手試試用OpenVINO 和 LCMs 吧！

關(guān)于OpenVINO 開(kāi)源工具套件的詳細(xì)資料[6]，包括其中我們提供的三百多個(gè)經(jīng)驗(yàn)證并優(yōu)化的預(yù)訓(xùn)練模型的詳細(xì)資料。

除此之外，為了方便大家了解并快速掌握OpenVINO 的使用，我們還提供了一系列開(kāi)源Jupyter notebook demo，運(yùn)行這些notebook，就能快速了解在不同場(chǎng)景下如何利用OpenVINO實(shí)現(xiàn)一系列、包括計(jì)算機(jī)視覺(jué)、語(yǔ)音及自然語(yǔ)言處理任務(wù)。

審核編輯：湯梓紅