一步到位：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）中利用图像潜在空间有效地增强了图像生成质量并减少了计算负载，使得图片急速生成成为了可能。（有关所提出方法和模型的更多详细信息，可以在项目页面Latent Consistency Models: Synthesizing High-Resolution Images with Few-step Inference 、论文https://arxiv.org/abs/2310.04378 和原始仓库GitHub - luosiallen/latent-consistency-model: Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference 中找到。）

如此充满魔力的LCMs文生图模型，我们的OpenVINO™当然也可以对它进行完全的优化、压缩以及推理加速、快速部署的支持。接下来，就让我们通过我们常用的OpenVINO Notebooks仓库中关于LCMs模型的Jupyter Notebook代码和拆解，来进一步了解具体步骤吧（Jupyter notebook代码地址https://github.com/openvinotoolkit/openvino_notebooks/tree/main/notebooks/263-latent-consistency-models-image-generation ）。

第一步：安装相应工具包、加载模型并转换为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三个模型。其中
1. Text Encoder 模型负责从文本text prompt 创建图片生成的条件
2. U-Net 模型负责初步降噪潜在图像表示
3. 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解码器解码为最终的图像输出。

扫描二维码关注公众号，回复： 17231190 查看本文章

定义关于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、

以及锐炫ARC A770m独立显卡运行了模型推理，生成图片真的是在眨眼之间！

当然，为了方便各位开发者的使用，我们的notebook 代码中还为大家设计了基于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)

小结：

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

关于英特尔OpenVINOTM开源工具套件的详细资料，包括其中我们提供的三百多个经验证并优化的预训练模型的详细资料，请您点击https://www.intel.com/content/www/us/en/developer/tools/openvino-toolkit/overview.html

除此之外，为了方便大家了解并快速掌握OpenVINOTM的使用，我们还提供了一系列开源的Jupyter notebook demo。运行这些notebook，就能快速了解在不同场景下如何利用OpenVINOTM实现一系列、包括计算机视觉、语音及自然语言处理任务。OpenVINOTM notebooks的资源可以在GitHub这里下载安装：https://github.com/openvinotoolkit/openvino_notebooks 。