OpenVINO赋能潜在一致性模型(LCMs)的高速图像生成

 

作者：武卓 英特尔 AI 布道师

 

提起生成式 AI，从去年开始爆火至今，相信无论是不是 AI 领域的开发者，对这个概念都是毫不陌生了。说到生成式 AI 领域的典型应用场景和模型，相信大家最先想到的就是文生图这个场景，以及它背后一系列的潜在扩散模型（Latent Diffusion models, LDMs）。借由这些模型的强大能力，只需要输入一段文字，人人都可化身设计师、插画师，进行精美图片的创作。但是，调用文生图模型的时候，若是选择调用云端运行的 API ，可能需要付费的同时、还有可能面临排队等待等问题。如果选择在本地机器上运行，那么对于机器的算力和内存就有较高的要求，同时也需要进行长时间的等待，毕竟这些模型往往需要迭代几十次之后才能生成一幅比较精美的图片。

最近一个称之为 LCMs（Latent Consistency Models）的模型横空出世，让文生图模型的图片急速生成成为了可能。受到一致性模型（CM）的启发，潜在一致性模型（LCMs），可以在任何预训练的潜在扩散模型上进行快速推断、步骤最少，包括稳定扩散模型（Stable Diffusion）。一致性模型是一种新的生成模型家族，可以实现一步或少量步骤生成。其核心思想是学习 PF-ODE（常微分方程概率流的轨迹）解的函数。通过学习保持 ODE 轨迹上点的一致性的一致性映射，这些模型允许进行单步生成，消除了计算密集的迭代的需求。然而，CM 受限于像素空间图像生成任务，因此不适用于合成高分辨率图像。LCMs 在图像潜在空间采用一致性模型来生成高分辨率图像。将引导反向扩散过程视为解决 PF-ODE 的过程。LCMs 旨在直接预测潜在空间中这种 ODE 的解，减少了大量迭代的需求，实现了快速、高保真度的采样。在大规模扩散模型（如 Stable Diffusion）中利用图像潜在空间有效地增强了图像生成质量并减少了计算负载，使得图片急速生成成为了可能。

有关所提出方法和模型的更多详细信息，可以在项目页面[1]、论文[2] 和原始仓库[3] 中找到。

如此充满魔力的 LCMs 文生图模型，我们的 OpenVINO 当然也可以对它进行完全的优化、压缩以及推理加速、快速部署的支持。接下来，就让我们通过我们常用的 OpenVINO Notebooks 仓库中关于 LCMs 模型的 Jupyter Notebook 代码[4] 和拆解，来进一步了解具体步骤吧。

 

第一步： 

安装相应工具包、加载模型

并转换为 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

 

左滑查看更多

 

模型下载

与传统的 Stable Diffusion 流水线类似， LCMs 模型内部也包含了 Text Encoder, U-Net, VAE Decoder 三个模型。

 

Text Encoder 模型

负责从文本 text prompt 创建图片生成的条件

 

U-Net 模型

负责初步降噪潜在图像表示

 

Autoencoder (VAE) Decoder 

用于将潜在空间解码为最终的图片

因此分别需要对这三个模型进行下载，部分代码如下：

 

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)

 

左滑查看更多

 

模型转换

包括将以上三个模型的转换为 OpenVINO IR 格式。

 

第二步： 

准备基于 OpenVINO 

的推理流水线

流水线如下图所示。

整个流水线利用一个潜在图像表示和一个文本提示，通过 CLIP 的文本编码器将文本提示转化为文本嵌入作为输入。初始的潜在图像表示是使用随机噪声生成器生成的。与原始的 Stable Diffusion 流程不同，LCMs 还使用引导尺度来获取时间步骤条件嵌入作为扩散过程的输入，而在 Stable Diffusion 中，它用于缩放输出的潜在表示。

接下来，U-Net 迭代地去噪随机潜在图像表示，同时以文本嵌入为条件。U-Net 的输出是噪声残差，通过调度算法用于计算去噪后的潜在图像表示。LCMs 引入了自己的调度算法，扩展了去噪扩散概率模型（DDPMs）中引入的非马尔可夫引导。去噪过程会重复多次（注意：原始 SD 流程中默认为 50 次，但对于 LCM，只需要 2-8 次左右的小步骤即可！），以逐步获取更好的潜在图像表示。完成后，潜在图像表示将由变 VAE 解码器解码为最终的图像输出。

定义关于 LCMs 推理流水线的类，部分代码如下：

 

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):

 

左滑查看更多

 

第三步：

配置推理流水线

首先，创建 OpenVINO 模型的实例，并使用选定的设备对其进行编译。从下拉列表中选择设备，以便使用 OpenVINO 运行推理。

 

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）和调度器（scheduler）也是流水线的重要组成部分。该流水线还可以使用安全检查器，该过滤器用于检测相应生成的图像是否包含“不安全工作”（nsfw）内容。nsfw 内容检测过程需要使用 CLIP 模型获得图像嵌入，因此需要在流水线中添加额外的特征提取器组件。我们重用原始 LCMs 流水线中的标记器、特征提取器、调度器和安全检查器。

 

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

 

左滑查看更多

在我本地的机器上，我分别使用了我的英特尔 酷睿 第 12 代 CPU、以及英特尔锐炫 显卡 A770m 运行了模型推理，生成图片真的是在眨眼之间！

当然，为了方便各位开发者的使用，我们的notebook 代码[5] 中还为大家设计了基于 Gradio 的、更加用户优化的界面。

 

第五步：

利用 NNCF 对模型进行量化压缩

另外，如果大家对于模型还有进一步的尺寸压缩以及内存占用压缩的需求，我们的 notebook 中也同样提供了基于 NNCF 进行量化压缩的代码示例。

量化压缩的过程分为以下三个步骤：

为量化创建一个校准数据集

运行 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)

 

左滑查看更多

运行效果如下：

根据同样的文本 text prompt，检测一下量化后模型的生成效果：

 

当然推理时间上的表现有何提升呢，我们也欢迎大家利用如下的代码在自己的机器上实测一下：

 

%%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)

 

左滑查看更多

 

小结

整个的步骤就是这样！现在就开始跟着我们提供的代码和步骤，动手试试用 OpenVINO 和 LCMs 吧！

关于 OpenVINO 开源工具套件的详细资料[6]，包括其中我们提供的三百多个经验证并优化的预训练模型的详细资料。

除此之外，为了方便大家了解并快速掌握OpenVINO 的使用，我们还提供了一系列开源 Jupyter notebook demo，运行这些 notebook，就能快速了解在不同场景下如何利用OpenVINO 实现一系列、包括计算机视觉、语音及自然语言处理任务。

　　审核编辑：汤梓红