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

import ipywidgets as widgets 

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

device

from optimum.intel.openvino import OVStableDiffusionXLPipeline 

text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value)

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

del text2image_pipe gc.collect();

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("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True).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)

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

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

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

%%skip not $to_quantize.value 

import time 

validation_size = 7 
calibration_dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True) 
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}")