SDXL-turbo と OpenVINO を使用したシングルステップの画像生成

この Jupyter ノートブックは、ローカルへのインストール後にのみ起動できます。

SDXL-Turbo は、単一のネットワーク評価でテキストプロンプトから写実的な画像を合成できる、高速な生成テキストから画像へのモデルです。SDXL-Turbo は、リアルタイム合成用にトレーニングされた SDXL 1.0 の抽出バージョンです。SDXL Turbo は、Adversarial Diffusion Distillation (ADD) と呼ばれる新しい抽出技術に基づいており、これによりモデルは単一のステップで画像出力を合成し、高いサンプリング忠実度を維持しながらリアルタイムのテキストから画像への出力を生成できます。この抽出アプローチの詳細については、技術レポートをご覧ください。モデルの詳細については、Stability AI のブログ投稿をご覧ください。

以前、次のノートブックで OpenVINO を使用して Stable Diffusion XL モデルを起動する方法についてすでに説明しましたが、このチュートリアルでは SDXL-turbo バージョンに焦点を当てます。さらに、画像のデコード速度を向上させるために、SDXL 生成プロセスのリアルタイム・プレビューに役立つ Tiny Autoencoder を使用します。

Hugging Face Diffusers ライブラリーの事前トレーニング済みモデルを使用します。ユーザー・エクスペリエンスを簡素化するために、Hugging Face Optimum Intel ライブラリーを使用してモデルを OpenVINO™ IR 形式に変換します。

目次

• 必要条件

• モデルを OpenVINO 形式に変換

• テキストから画像への生成

  • テキストから画像を生成する推論デバイスを選択

• Image-to-Image 生成

• 量子化

  • キャリブレーション・データセットの準備

  • 量子化を実行

    • UNet ファイルサイズを比較

  • FP16 モデルと INT8 モデルの推論時間を比較

• インタラクティブなデモ

必要条件

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

モデルを OpenVINO 形式に変換

sdxl-turbo は、Hugging Face Hub からダウンロードできます。OpenVINO 中間表現 (IR) 形式にエクスポートするには、optimal-cli インターフェイスを使用します。

モデルを変換するための Optimum CLI インターフェイスは、OpenVINO へのエクスポートをサポートします (optimum-intel バージョン 1.12 以降でサポートされます)。
一般的なコマンド形式:

optimum-cli export openvino --model <model_id_or_path> --task <task> <output_dir>

ここで、task はモデルをエクスポートするタスクです。指定されていない場合は、モデルに基づいてタスクが自動的に推論されます。利用可能なタスクはモデルによって異なります。SDXL の場合は stable-diffusion-xl を選択する必要があります。

タスクとモデルクラス間のマッピングについては、Optimum TaskManager のドキュメントを参照してください。

さらに、圧縮モデルを FP16 にする場合は --fp16 の重みを指定し、圧縮モデルを INT8 にする場合は --int8 の重みを指定することもできます。INT8 の場合は nncf をインストールする必要があることに注意してください。

サポートされている引数の完全なリストは --help で参照できます。詳細と使用例については、Optimum ドキュメントを確認してください。

Tiny Autoencoder の場合、ov.convert_model 関数を使用して ov.Model を取得し、ov.save_model を使用して保存します。モデルは、パイプラインで別々に使用される 2 つの部分で構成されています。1 つは、画像から画像への生成タスクで潜在空間に入力画像をエンコードする vae_encoder で、もう 1 つは拡散結果を画像形式にデコードする vae_decoder です。

from pathlib import Path

model_dir = Path("./model")
sdxl_model_id = "stabilityai/sdxl-turbo"
tae_id = "madebyollin/taesdxl"
skip_convert_model = model_dir.exists()

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(model_id, output_path):
    tiny_vae = AutoencoderTiny.from_pretrained(model_id)
    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")


if not skip_convert_model:
    !optimum-cli export openvino --model $sdxl_model_id --task stable-diffusion-xl $model_dir --fp16
    convert_tiny_vae(tae_id, model_dir)

テキストから画像への生成

テキストから画像への生成により、テキストの説明を使用して画像を作成できます。画像の生成を開始するには、まずモデルを読み込む必要があります。OpenVINO モデルをロードして OpenVINO ランタイムで推論を実行するには、Diffusers の StableDiffusionXLPipeline を Optimum の OVStableDiffusionXLPipeline に置き換える必要があります。パイプラインの初期化は from_pretrained メソッドを使用して開始され、OpenVINO モデルを含むディレクトリーを渡す必要があります。さらに、推論デバイスを指定することもできます。

テキストから画像を生成する推論デバイスを選択

import ipywidgets as widgets

core = ov.Core()

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

device

Dropdown(description='Device:', index=1, options=('CPU', 'AUTO'), value='AUTO')

from optimum.intel.openvino import OVStableDiffusionXLPipeline
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)

INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino

/home/ea/work/genai_env/lib/python3.8/site-packages/torch/cuda/__init__.py:138: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 11080). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)
  return torch._C._cuda_getDeviceCount() > 0
No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
Compiling the vae_decoder to AUTO ...
Compiling the unet to AUTO ...
Compiling the text_encoder to AUTO ...
Compiling the text_encoder_2 to AUTO ...
Compiling the vae_encoder to AUTO ...

パイプライン・インターフェースは、元の StableDiffusionXLPipeline に似ています。テキストプロンプトを提供する必要があります。デフォルトのステップ数は 50 ですが、sdxl-turbo では 1 ステップのみ必要です。モデルカードに記載されている情報によると、モデルはネガティブ・プロンプトとガイダンススケールを使用していないため、このパラメーターは次のように無効にする必要があります: guidance_scale = 0。

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

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

del text2image_pipe
gc.collect();

Image-to-Image 生成

画像から画像への生成により、テキストの説明で提供される特性に合わせて画像を変換できます。すでに変換されたモデルを再利用して、Image2Image 生成パイプラインを実行することができます。それには、OVStableDiffusionXLPipeline を OVStableDiffusionXLImage2ImagePipeline に置き換える必要があります。

from optimum.intel import OVStableDiffusionXLImg2ImgPipeline

image2image_pipe = OVStableDiffusionXLImg2ImgPipeline.from_pretrained(model_dir, device=device.value)

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

photo_prompt = "a cute cat with bow tie"

strength パラメーターは、画像から画像への生成パイプラインにとって重要です。これは 0.0 から 1.0 の間の値で、入力画像に追加されるノイズの量を制御します。値が 1.0 に近づくと、さまざまなバリエーションが可能になりますが、入力と意味的に一致しない画像も生成されます。0 に近づくと、追加されるノイズが少なくなり、ターゲット画像はソース画像の内容を保持します。強度は、ノイズの数だけでなく、生成ステップの数にも影響します。画像から画像への生成パイプラインにおけるノイズ除去の反復回数は、int(num_inference_steps * strength) として計算されます。sdxl-turbo では、正しい結果を生成するために num_inference_steps と strength の選択に注意し、強度乗算を適用した後、パイプラインで使用されるステップ数が >=1 であることを確認する必要があります。例えば、以下の例では、num_inference_steps=2 および stength=0.5 を使用し、最終的にパイプラインで 0.5 * 2.0 = 1 ステップを取得します。

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

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

del image2image_pipe
gc.collect();

量子化

NNCF は、量子化レイヤーをモデルグラフに追加し、トレーニング・データセットのサブセットを使用してこれらの追加の量子化レイヤーのパラメーターを初期化することで、トレーニング後の量子化を可能にします。量子化操作は FP32/FP16 ではなく INT8 で実行されるため、モデル推論が高速化されます。

SDXL-Turbo モデル構造によれば、UNet モデルはパイプライン全体の実行時間の大部分を占めます。ここでは、NNCF を使用して UNet 部分を最適化し、計算コストを削減してパイプラインを高速化する方法を説明します。SDXL パイプラインの残りのパーツを量子化しても、推論パフォーマンスは大幅に向上せず、精度が大幅に低下する可能性があります。

最適化プロセスには次の手順が含まれます。

1. 量子化用のキャリブレーション・データセットを作成します。

2. nncf.quantize() を実行して、量子化されたモデルを取得します。

3. openvino.save_model() 関数を使用して INT8 モデルを保存します。

モデルの推論速度を向上させるため量子化を実行するかどうかを以下で選択してください。

to_quantize = widgets.Checkbox(
    value=True,
    description='Quantization',
    disabled=False,
)

to_quantize

Checkbox(value=True, description='Quantization')

import sys
sys.path.append("../utils")

int8_pipe = None

if to_quantize.value and "GPU" in device.value:
    to_quantize.value = False

%load_ext skip_kernel_extension

キャリブレーション・データセットの準備

Hugging Face の検証 conceptual_captions データセットの一部をキャリブレーション・データとして使用します。キャリブレーション用の中間モデル入力を収集するには、CompiledModel をカスタマイズする必要があります。

UNET_INT8_OV_PATH = model_dir / "optimized_unet" / "openvino_model.xml"

def disable_progress_bar(pipeline, disable=True):
    if not hasattr(pipeline, "_progress_bar_config"):
        pipeline._progress_bar_config = {'disable': disable}
    else:
        pipeline._progress_bar_config['disable'] = disable

%%skip not $to_quantize.value

import datasets
import numpy as np
from tqdm.notebook 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)
    disable_progress_bar(pipe)

    # 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
    disable_progress_bar(pipe, disable=False)
    pipe.unet.request = original_unet
    return calibration_dataset

%%skip not $to_quantize.value

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

量子化を実行

事前トレーニング済みの変換済み OpenVINO モデルから量子化モデルを作成します。最初と最後の畳み込みレイヤーの量子化は、生成結果に影響します。FP16 精度で精度に敏感な畳み込みレイヤーを維持するには、IgnoredScope を使用することをお勧めします。

注: 量子化は時間とメモリーを消費する操作です。以下の量子化コードの実行には時間がかかる場合があります。

%%skip not $to_quantize.value

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 で予測を確認してみましょう。

%%skip not $to_quantize.value

from IPython.display import display

int8_text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value, 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 AUTO ...
Compiling the text_encoder_2 to AUTO ...

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

Compiling the unet to AUTO ...
Compiling the vae_decoder to AUTO ...

%%skip not $to_quantize.value

int8_image2image_pipe = OVStableDiffusionXLImg2ImgPipeline.from_pretrained(model_dir, device=device.value, 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 AUTO ...
Compiling the text_encoder_2 to AUTO ...
Compiling the vae_encoder to AUTO ...

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

Compiling the unet to AUTO ...
Compiling the vae_decoder to AUTO ...

UNet ファイルサイズを比較

%%skip not $to_quantize.value

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: 2513541.44 KB
Model compression rate: 1.995

FP16 および INT8 パイプラインの推論パフォーマンスを測定するには、キャリブレーション・サブセットの推論時間の中央値を使用します。

注: 最も正確なパフォーマンス推定を行うには、他のアプリケーションを閉じた後、ターミナル/コマンドプロンプトで benchmark_app を実行することを推奨します。

%%skip not $to_quantize.value

import time

validation_size = 7
calibration_dataset = datasets.load_dataset("conceptual_captions", split="train")
validation_data = []
for batch in calibration_dataset:
    prompt = batch["caption"]
    validation_data.append(prompt)

def calculate_inference_time(pipe, dataset):
    inference_time = []
    disable_progress_bar(pipe)

    for idx, prompt in enumerate(dataset):
        start = time.perf_counter()
        image = 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)
        if idx >= validation_size:
            break
    disable_progress_bar(pipe, disable=False)
    return np.median(inference_time)

%%skip not $to_quantize.value

int8_latency = calculate_inference_time(int8_text2image_pipe, validation_data)
text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)
fp_latency = calculate_inference_time(text2image_pipe, validation_data)
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}")

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

FP16 pipeline latency: 1.391
INT8 pipeline latency: 0.781
Text-to-Image generation speed up: 1.780

インタラクティブなデモ

独自のテキスト記述を使用してモデル作業を確認できるようになりました。テキストボックスにテキストプロンプトを入力し、[Run] ボタンで生成を開始します。さらに、追加のパラメーターを使用して生成を制御することもできます。

• シード - 初期化用のランダムシード
• ステップ - 生成ステップの数
• 高さと幅 - 生成された画像のサイズ

画像サイズを大きくすると、正確な結果を得るのに必要な手順数が増える場合があることに注意してください。104x1024 解像度の画像生成は 4 つのステップで実行することを推奨します。

インタラクティブなデモを起動するため量子化モデルを使用するかどうか以下で選択してください。

quantized_model_present = UNET_INT8_OV_PATH.exists()

use_quantized_model = widgets.Checkbox(
    value=True if quantized_model_present else False,
    description='Use quantized model',
    disabled=False,
)

use_quantized_model

Checkbox(value=True, description='Use quantized model')

import gradio as gr

text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)
if use_quantized_model.value:
    if not quantized_model_present:
        raise RuntimeError("Quantized model not found.")
    text2image_pipe.unet.model = core.read_model(UNET_INT8_OV_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)" if use_quantized_model.value else "Result (Original)", 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])

# if you are launching remotely, specify server_name and server_port
# demo.launch(server_name='your server name', server_port='server port in int')
# Read more in the docs: https://gradio.app/docs/
# if you want create public link for sharing demo, please add share=True
try:
    demo.launch(debug=False)
except Exception:
    demo.launch(share=True, debug=False)