一步到位：利用OpenVINO實現LCM高速影像生成

Posted By 武卓 on 11 月 17, 2023 in Edge AI教學文, Edge AI開發地圖, LCM, OpenVINO, OpenVINO專欄, OpenVINO教學文, Stable Diffusion, 教學文, 開發套件 | 0 comments

提起從去年開始爆紅至今的生成式AI，相信無論是不是AI領域的開發者，都對這個概念毫不陌生。而說到生成式AI領域的典型應用場景和模型，大家最先想到的應該就是文生圖，以及它背後一系列的潛在擴散模型(Latent Diffusion model，LDM)；借助這些模型的強大能力，只需要輸入一段文字，人人都可化身設計師、插畫師，進行精美圖片的創作。

但是調用文生圖模型的時候，若選擇雲端運作的API，不僅可能需要付費、還有可能面臨排隊等待等問題。如果選擇在本地機器上運作，那麼對於機器的算力和記憶體就有較高的要求，同時也需要進行長時間的等待，畢竟這些模型往往需要迭代幾十次之後才能生成一幅比較精美的圖片。

而最近橫空出世的潛在一致性模型(Latent Consistency Model，LCM)，讓文生圖模型的圖片急速生成變成可能。受到一致性模型(CM)的啟發，LCM可以在任何預訓練的潛在擴散模型上進行快速推斷、步驟最少，包括穩定擴散模型(Stable Diffusion)。

一致性模型是一種新的生成模型家族，可以實現一步或少量步驟生成。其核心思想是學習PF-ODE (常微分方程式概率流的軌跡)解的函數。透過學習保持ODE軌跡上點的一致性的一致性映射，這些模型允許進行單步生成，消除了運算密集的迭代需求。

然而，CM受限於畫素空間影像生成任務，因此不適用於合成高解析度影像。LCM在影像潛在空間採用一致性模型來生成高解析度影像。將引導反向擴散過程視為解決PF-ODE的過程。LCM旨在直接預測潛在空間中這種ODE的解，減少了大量迭代的需求，實現了快速、高傳真度的採樣。

在大規模擴散模型(如Stable Diffusion)中利用影像潛在空間有效地增強了影像生成品質並減少了運算負載，使得圖片急速生成變成可能(有關所提出方法和模型的更多詳細資訊，可以在專案頁面、論文和原始碼倉庫中找到)。

如此充滿魔力的LCM文生圖模型，我們的OpenVINO當然也可以對它進行完全的最佳化、壓縮以及推理加速、快速部署的支援。接下來，就讓我們透過常用的OpenVINO Notebooks資源庫中關於LCM的Jupyter Notebook程式碼(點此連結)和拆解，來進一步瞭解具體步驟吧！

OpenVINO^™更多的新功能陸續問世！新版本將新增對新硬體的支援，並支援更多生成式AI模型。在2024年的新開始，讓我們一起透過免費線上講座「掌握OpenVINO在生成式AI最新功能, 並實現Model Server創新應用」來探索這套越來越強大的工具，讓您的專案開發更如虎添翼，名額有限快來報名！

第一步： 安裝相應工具套件、載入模型並轉換為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流水線類似，LCM內部也包含了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流程不同，LCM還使用引導尺度來獲取時間步驟條件嵌入作為擴散過程的輸入，而在Stable Diffusion中，它用於縮放輸出的潛在表示。

接下來，U-Net迭代對隨機潛在影像表示去雜訊，同時以文字嵌入為條件。U-Net的輸出是雜訊殘差，透過調度演算法用於計算去雜訊後的潛在影像表示。LCMs導入了自己的調度演算法，擴展了去雜訊擴散概率模型(DDPM)中導入的非馬可夫(non-Markovian)引導。

去雜訊過程會重複多次(注意：原始SD流程中預設為50次，但對於LCM只需要2-8次左右的小步驟即可！)，以逐步獲得更好的潛在影像表示。完成後，潛在影像表示將由變分自解碼器(VAE)解碼為最終的影像輸入。

定義關於LCM推論流水線的類別，部分程式碼如下：

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 + ,
    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代酷睿(Core) CPU、以及銳炫(ARC) A770m獨立顯卡執行了模型推論，生成圖片真的是在眨眼之間！