作者：英特尔创新大使 卢雨畋

本文基于 13th Gen Intel(R) Core(TM) i5-13490F 型号 CPU 验证,对于量化后模型，你只需要在 16G 的笔记本电脑上就可体验生成过程（最佳体验为32G内存）。
SDXL-Turbo是一个快速的生成式文本到图像模型，可以通过单次网络评估从文本提示中合成逼真的图像。SDXL-Turbo采用了一种称为Adversarial Diffusion Distillation (ADD)的新型训练方法（详见技术报告），该方法可以在1到4个步骤中对大规模基础图像扩散模型进行采样，并保持高质量的图像。通过最新版本（2023.2）OpenVINO™ 工具套件的强大推理能力及 NNCF 的高效神经网络压缩能力，我们能够在两秒内实现 SDXL-Turbo 图像的高速、高质量生成。

环境安装

在开始之前，我们需要安装所有环境依赖：

%pip install --extra-index-url https://download.pytorch.org/whl/cpu \
torch transformers diffusers nncf optimum-intel gradio openvino==2023.2.0 onnx "git+https://github.com/huggingface/optimum-intel.git"

下载、转换模型

首先我们要把 huggingface 下载的原始模型转化为 OpenVINO IR，以便后续的 NNCF 工具链进行量化工作。转换完成后你将得到对应的 text_encode、unet、vae模型。

from pathlib import Path
model_dir = Path("./sdxl_vino_model")
sdxl_model_id = "stabilityai/sdxl-turbo"
skip_convert_model = model_dir.exists()

import os
if not skip_convert_model:
    # 设置下载路径到当前文件夹，并加速下载
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    os.system(f'optimum-cli export openvino --model {sdxl_model_id} --task stable-diffusion-xl {model_dir} --fp16')
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
tae_id = "madebyollin/taesdxl"
save_path = './taesdxl'
os.system(f'huggingface-cli download --resume-download {tae_id} --local-dir {save_path}')

import torch
import openvino as ov
from diffusers import AutoencoderTiny
import gc

class VAEEncoder(torch.nn.Module):
    def __init__(self, vae):
        super().__init__()
        self.vae = vae

    def forward(self, sample):
        return self.vae.encode(sample)
   
class VAEDecoder(torch.nn.Module):
    def __init__(self, vae):
        super().__init__()
        self.vae = vae

    def forward(self, latent_sample):
        return self.vae.decode(latent_sample)

def convert_tiny_vae(save_path, output_path):
    tiny_vae = AutoencoderTiny.from_pretrained(save_path)
    tiny_vae.eval()
    vae_encoder = VAEEncoder(tiny_vae)
    ov_model = ov.convert_model(vae_encoder, example_input=torch.zeros((1,3,512,512)))
    ov.save_model(ov_model, output_path / "vae_encoder/openvino_model.xml")
    tiny_vae.save_config(output_path / "vae_encoder")
    vae_decoder = VAEDecoder(tiny_vae)
    ov_model = ov.convert_model(vae_decoder, example_input=torch.zeros((1,4,64,64)))
    ov.save_model(ov_model, output_path / "vae_decoder/openvino_model.xml")
    tiny_vae.save_config(output_path / "vae_decoder")    

convert_tiny_vae(save_path, model_dir)

从文本到图像生成

现在，我们就可以进行文本到图像的生成了，我们使用优化后的 openvino pipeline 加载转换后的模型文件并推理；只需要指定一个文本输入，就可以生成我们想要的图像结果。

from optimum.intel.openvino import OVStableDiffusionXLPipeline
device='AUTO'  # 这里直接指定AUTO，可以写成CPU
model_dir = "./sdxl_vino_model"
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device)

import numpy as np
prompt = "cute cat"
image = text2image_pipe(prompt, num_inference_steps=1, height=512, width=512, guidance_scale=0.0, generator=np.random.RandomState(987)).images[0]
image.save("cat.png")
image

# 清除资源占用
import gc
del text2image_pipe
gc.collect()

从图片到图片生成

我们还可以实现从图片到图片的扩散模型生成，将刚才产出的文生图图片进行二次图像生成即可。

from optimum.intel import OVStableDiffusionXLImg2ImgPipeline
model_dir = "./sdxl_vino_model"
device='AUTO'  # 'CPU'
image2image_pipe = OVStableDiffusionXLImg2ImgPipeline.from_pretrained(model_dir, device=device)

    Compiling the vae_decoder to AUTO ...
    Compiling the unet to AUTO ...
    Compiling the vae_encoder to AUTO ...
    Compiling the text_encoder_2 to AUTO ...
    Compiling the text_encoder to AUTO ...

photo_prompt = "a cute cat with bow tie"
photo_image = image2image_pipe(photo_prompt, image=image, num_inference_steps=2, generator=np.random.RandomState(511), guidance_scale=0.0, strength=0.5).images[0]
photo_image.save("cat_tie.png")
photo_image

量化

NNCF(Neural Network Compression Framework)是一款神经网络压缩框架，通过对OpenVINO IR格式模型的压缩与量化巍以便更好的提升模型在 Intel 设备上部署的推理性能。
[NNCF](https://github.com/openvinotoolkit/nncf/) 通过在模型图中添加量化层，并使用训练数据集的子集来微调这些额外的量化层的参数，实现了后训练量化。量化后的权重结果将是 INT8 而不是 FP32/FP16，从而加快了模型的推理速度。
根据 SDXL-Turbo Model 的结构，UNet 模型占据了整个流水线执行时间的重要部分。现在我们将展示如何使用  [NNCF](https://github.com/openvinotoolkit/nncf/)  对 UNet 部分进行优化，以减少计算成本并加快流水线速度。至于其余部分不需要量化，因为并不能显著提高推理性能，但可能会导致准确性的大幅降低。

量化过程包含以下步骤：
- 为量化创建一个校准数据集。
- 运行 nncf.quantize() 来获取量化模型。
- 使用 openvino.save_model() 函数保存 INT8 模型。
注：由于量化需要一定的硬件资源（64G以上的内存），之后我直接附上了量化后的模型，你可以直接下载使用。

from pathlib import Path
import openvino as ov
from optimum.intel.openvino import OVStableDiffusionXLPipeline
import os

core = ov.Core()
model_dir = Path("./sdxl_vino_model")
UNET_INT8_OV_PATH = model_dir / "optimized_unet" / "openvino_model.xml"

import datasets
import numpy as np
from tqdm import tqdm
from transformers import set_seed
from typing import Any, Dict, List

set_seed(1)

class CompiledModelDecorator(ov.CompiledModel):
    def __init__(self, compiled_model: ov.CompiledModel, data_cache: List[Any] = None):
        super().__init__(compiled_model)
        self.data_cache = data_cache if data_cache else []

    def __call__(self, *args, **kwargs):
        self.data_cache.append(*args)
        return super().__call__(*args, **kwargs)

def collect_calibration_data(pipe, subset_size: int) -> List[Dict]:
    original_unet = pipe.unet.request
    pipe.unet.request = CompiledModelDecorator(original_unet)
    dataset = datasets.load_dataset("conceptual_captions", split="train").shuffle(seed=42)

    # Run inference for data collection
    pbar = tqdm(total=subset_size)
    diff = 0
    for batch in dataset:
        prompt = batch["caption"]
        if len(prompt) > pipe.tokenizer.model_max_length:
            continue
        _ = pipe(
            prompt,
            num_inference_steps=1,
            height=512,
            width=512,
            guidance_scale=0.0,
            generator=np.random.RandomState(987)
        )
        collected_subset_size = len(pipe.unet.request.data_cache)
        if collected_subset_size >= subset_size:
            pbar.update(subset_size - pbar.n)
            break
        pbar.update(collected_subset_size - diff)
        diff = collected_subset_size

    calibration_dataset = pipe.unet.request.data_cache
    pipe.unet.request = original_unet
    return calibration_dataset

if not UNET_INT8_OV_PATH.exists():
    text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir)
    unet_calibration_data = collect_calibration_data(text2image_pipe, subset_size=200)

import nncf
from nncf.scopes import IgnoredScope

UNET_OV_PATH = model_dir / "unet" / "openvino_model.xml"
if not UNET_INT8_OV_PATH.exists():
    unet = core.read_model(UNET_OV_PATH)
    quantized_unet = nncf.quantize(
        model=unet,
        model_type=nncf.ModelType.TRANSFORMER,
        calibration_dataset=nncf.Dataset(unet_calibration_data),
        ignored_scope=IgnoredScope(
            names=[
                "__module.model.conv_in/aten::_convolution/Convolution",
                "__module.model.up_blocks.2.resnets.2.conv_shortcut/aten::_convolution/Convolution",
                "__module.model.conv_out/aten::_convolution/Convolution"
            ],
        ),
    )
    ov.save_model(quantized_unet, UNET_INT8_OV_PATH)

运行量化后模型

由于量化 unet 的过程需要的内存可能比较大，且耗时较长，我提前导出了量化后 unet 模型，此处给出下载地址：
链接: https://pan.baidu.com/s/1WMAsgFFkKKp-EAS6M1wK1g 提取码: psta
下载后解压到目标文件夹 `sdxl_vino_model` 即可运行量化后的 int8 unet 模型。
从文本到图像生成

from pathlib import Path
import openvino as ov
from optimum.intel.openvino import OVStableDiffusionXLPipeline
import numpy as np

core = ov.Core()
model_dir = Path("./sdxl_vino_model")
UNET_INT8_OV_PATH = model_dir / "optimized_unet" / "openvino_model.xml"
int8_text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, compile=False)
int8_text2image_pipe.unet.model = core.read_model(UNET_INT8_OV_PATH)
int8_text2image_pipe.unet.request = None

prompt = "cute cat"
image = int8_text2image_pipe(prompt, num_inference_steps=1, height=512, width=512, guidance_scale=0.0, generator=np.random.RandomState(987)).images[0]
display(image)

 Compiling the text_encoder to CPU ...
    Compiling the text_encoder_2 to CPU ...


      0%|          | 0/1 [00:00<?, ?it/s]

    Compiling the unet to CPU ...
    Compiling the vae_decoder to CPU ...

从图片到图片生成

from optimum.intel import OVStableDiffusionXLImg2ImgPipeline
int8_image2image_pipe = OVStableDiffusionXLImg2ImgPipeline.from_pretrained(model_dir, compile=False)
int8_image2image_pipe.unet.model = core.read_model(UNET_INT8_OV_PATH)
int8_image2image_pipe.unet.request = None

photo_prompt = "a cute cat with bow tie"
photo_image = int8_image2image_pipe(photo_prompt, image=image, num_inference_steps=2, generator=np.random.RandomState(511), guidance_scale=0.0, strength=0.5).images[0]
display(photo_image)


    Compiling the text_encoder to CPU ...
    Compiling the text_encoder_2 to CPU ...
    Compiling the vae_encoder to CPU ...


      0%|          | 0/1 [00:00<?, ?it/s]

    Compiling the unet to CPU ...
    Compiling the vae_decoder to CPU ...

我们可以对比量化后的 unet 模型大小减少，可以看到量化对模型大小的压缩是非常显著的

from pathlib import Path

model_dir = Path("./sdxl_vino_model")
UNET_OV_PATH = model_dir / "unet" / "openvino_model.xml"
UNET_INT8_OV_PATH = model_dir / "optimized_unet" / "openvino_model.xml"

fp16_ir_model_size = UNET_OV_PATH.with_suffix(".bin").stat().st_size / 1024
quantized_model_size = UNET_INT8_OV_PATH.with_suffix(".bin").stat().st_size / 1024

print(f"FP16 model size: {fp16_ir_model_size:.2f} KB")
print(f"INT8 model size: {quantized_model_size:.2f} KB")
print(f"Model compression rate: {fp16_ir_model_size / quantized_model_size:.3f}")

  FP16 model size: 5014578.27 KB
    INT8 model size: 2513501.39 KB
    Model compression rate: 1.995

运行下列代码可以对量化前后模型推理速度进行简单比较，我们可以发现速度几乎加速了一倍，NNCF 使我们在CPU上生成一张图的时间缩短到两秒之内：

FP16 pipeline latency: 3.148
INT8 pipeline latency: 1.558
Text-to-Image generation speed up: 2.020

import time
def calculate_inference_time(pipe):
    inference_time = []
    for prompt in ['cat']*10:
        start = time.perf_counter()
        _ = pipe(
            prompt,
            num_inference_steps=1,
            guidance_scale=0.0,
            generator=np.random.RandomState(23)
        ).images[0]
        end = time.perf_counter()
        delta = end - start
        inference_time.append(delta)
    return np.median(inference_time)

int8_latency = calculate_inference_time(int8_text2image_pipe)
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir)
fp_latency = calculate_inference_time(text2image_pipe)
print(f"FP16 pipeline latency: {fp_latency:.3f}")
print(f"INT8 pipeline latency: {int8_latency:.3f}")
print(f"Text-to-Image generation speed up: {fp_latency / int8_latency:.3f}")

可交互前端demo

最后，为了方便推理使用，这里附上了 gradio 前端运行 demo，你可以利用他轻松生成你想要生成的图像，并尝试不同组合。

import gradio as gr
from pathlib import Path
import openvino as ov
import numpy as np

core = ov.Core()
model_dir = Path("./sdxl_vino_model")

# 如果你只有量化前模型，请使用这个地址并注释 optimized_unet 地址：
# UNET_PATH = model_dir / "unet" / "openvino_model.xml"
UNET_PATH = model_dir / "optimized_unet" / "openvino_model.xml"

from optimum.intel.openvino import OVStableDiffusionXLPipeline
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir)
text2image_pipe.unet.model = core.read_model(UNET_PATH)
text2image_pipe.unet.request = core.compile_model(text2image_pipe.unet.model)

def generate_from_text(text, seed, num_steps, height, width):
    result = text2image_pipe(text, num_inference_steps=num_steps, guidance_scale=0.0, generator=np.random.RandomState(seed), height=height, width=width).images[0]
    return result

with gr.Blocks() as demo:
    with gr.Column():
        positive_input = gr.Textbox(label="Text prompt")
        with gr.Row():
            seed_input = gr.Number(precision=0, label="Seed", value=42, minimum=0)
            steps_input = gr.Slider(label="Steps", value=1, minimum=1, maximum=4, step=1)
            height_input = gr.Slider(label="Height", value=512, minimum=256, maximum=1024, step=32)
            width_input = gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=32)
            btn = gr.Button()
        out = gr.Image(label="Result (Quantized)" , type="pil", width=512)
        btn.click(generate_from_text, [positive_input, seed_input, steps_input, height_input, width_input], out)
        gr.Examples([
            ["cute cat", 999], 
            ["underwater world coral reef, colorful jellyfish, 35mm, cinematic lighting, shallow depth of field,  ultra quality, masterpiece, realistic", 89],
            ["a photo realistic happy white poodle dog playing in the grass, extremely detailed, high res, 8k, masterpiece, dynamic angle", 1569],
            ["Astronaut on Mars watching sunset, best quality, cinematic effects,", 65245],
            ["Black and white street photography of a rainy night in New York, reflections on wet pavement", 48199]
        ], [positive_input, seed_input])

try:
    demo.launch(debug=True)
except Exception:
    demo.launch(share=True, debug=True)

总结

利用最新版本的 OpenVINO™ 优化，我们可以很容易实现在家用设备上高效推理图像生成 AI 的能力，加速生成式 AI 在世纪场景下的落地应用；欢迎您与我们一同体验 OpenVINO™ 与 NNCF 在生成式 AI 场景上的强大威力。