Stable Diffusion v2 と OpenVINO™ によるテキストからの画像#

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base").to("cpu") 

# メモリー消費量を削減するために、パイプラインからすべてのコンポーネントを個別に取得 
text_encoder = pipe.text_encoder 
text_encoder.eval() 
unet = pipe.unet unet.eval() 
vae = pipe.vae 
vae.eval() 

conf = pipe.scheduler.config 

del pipe

from pathlib import Path 

sd2_1_model_dir = Path("sd2.1") 
sd2_1_model_dir.mkdir(exist_ok=True)

import gc 
import torch 
import openvino as ov 

TEXT_ENCODER_OV_PATH = sd2_1_model_dir / "text_encoder.xml" 

def cleanup_torchscript_cache(): 
    """ 
    Helper for removing cached model representation 
    """ 
    torch._C._jit_clear_class_registry() 
    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() 
    torch.jit._state._clear_class_state() 

def convert_encoder(text_encoder: torch.nn.Module, ir_path: Path): 
    """ 
    Convert Text Encoder model to IR.
    Function accepts pipeline, prepares example inputs for conversion 
    Parameters: 
        text_encoder (torch.nn.Module): text encoder PyTorch model 
        ir_path (Path): File for storing model 
    Returns:
        None 
    """ 
    if not ir_path.exists(): 
        input_ids = torch.ones((1, 77), dtype=torch.long) 
        # モデルを推論モードに切り替え 
        text_encoder.eval() 

        # メモリー消費量を減らすために勾配計算を無効化 
        with torch.no_grad():
            # エクスポート・モデル 
            ov_model = ov.convert_model( 
                text_encoder, # モデル・インスタンス 
                example_input=input_ids, # モデルトレースの入力例 
                input=([1, 77],), # 変換のための入力形状 
            ) 
            ov.save_model(ov_model, ir_path) 
            del ov_model 
            cleanup_torchscript_cache() 
        print("Text Encoder successfully converted to IR") 

if not TEXT_ENCODER_OV_PATH.exists(): 
    convert_encoder(text_encoder, TEXT_ENCODER_OV_PATH) 
else: 
    print(f"Text encoder will be loaded from {TEXT_ENCODER_OV_PATH}") 

del text_encoder 
gc.collect();

import numpy as np 

UNET_OV_PATH = sd2_1_model_dir / "unet.xml" 

def convert_unet( 
    unet: torch.nn.Module, 
    ir_path: Path, 
    num_channels: int = 4, 
    width: int = 64, 
    height: int = 64, 
): 
    """ 
    Convert Unet model to IR format.
    Function accepts pipeline, prepares example inputs for conversion 
    Parameters: 
        unet (torch.nn.Module): UNet PyTorch model 
        ir_path (Path): File for storing model 
        num_channels (int, optional, 4): number of input channels 
        width (int, optional, 64): input width 
        height (int, optional, 64): input height 
    Returns:
        None 
    """ 
    dtype_mapping = {torch.float32: ov.Type.f32, torch.float64: ov.Type.f64} 
    if not ir_path.exists():
        # 入力を準備 
        encoder_hidden_state = torch.ones((2, 77, 1024)) 
        latents_shape = (2, num_channels, width, height) 
        latents = torch.randn(latents_shape) 
        t = torch.from_numpy(np.array(1, dtype=np.float32)) 
        unet.eval() 
        dummy_inputs = (latents, t, encoder_hidden_state) 
        input_info = [] 
        for input_tensor in dummy_inputs: 
            shape = ov.PartialShape(tuple(input_tensor.shape)) 
            element_type = dtype_mapping[input_tensor.dtype] 
            input_info.append((shape, element_type)) 

        with torch.no_grad(): 
            ov_model = ov.convert_model(unet, example_input=dummy_inputs, input=input_info) 
        ov.save_model(ov_model, ir_path) 
        del ov_model 
        cleanup_torchscript_cache() 
        print("U-Net successfully converted to IR") 

if not UNET_OV_PATH.exists(): 
    convert_unet(unet, UNET_OV_PATH, width=96, height=96) 
    del unet 
    gc.collect() 
else: 
    del unet 
gc.collect();

VAE_ENCODER_OV_PATH = sd2_1_model_dir / "vae_encoder.xml" 

def convert_vae_encoder(vae: torch.nn.Module, ir_path: Path, width: int = 512, height: int = 512): 
    """ 
    Convert VAE model to IR format.
    VAE model, creates wrapper class for export only necessary for inference part, 
    prepares example inputs for onversion 
    Parameters: 
        vae (torch.nn.Module): VAE PyTorch model 
        ir_path (Path): File for storing model 
        width (int, optional, 512): input width 
        height (int, optional, 512): input height 
    Returns:
        None 
    """ 

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

        def forward(self, image): 
            return self.vae.encode(x=image)["latent_dist"].sample() 

        if not ir_path.exists(): 
            vae_encoder = VAEEncoderWrapper(vae) 
            vae_encoder.eval() 
            image = torch.zeros((1, 3, width, height)) 
            with torch.no_grad(): 
                ov_model = ov.convert_model(vae_encoder, example_input=image, input=([1, 3, width, height],)) 
            ov.save_model(ov_model, ir_path) 
            del ov_model 
            cleanup_torchscript_cache() 
            print("VAE encoder successfully converted to IR") 

def convert_vae_decoder(vae: torch.nn.Module, ir_path: Path, width: int = 64, height: int = 64): 
    """ 
    Convert VAE decoder model to IR format.
    Function accepts VAE model, creates wrapper class for export only necessary for inference part, 
    prepares example inputs for conversion, 
    Parameters: 
        vae (torch.nn.Module): VAE model 
        ir_path (Path): File for storing model 
        width (int, optional, 64): input width 
        height (int, optional, 64): input height 
    Returns:
        None 
    """ 

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

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

    if not ir_path.exists(): 
        vae_decoder = VAEDecoderWrapper(vae) 
        latents = torch.zeros((1, 4, width, height)) 

        vae_decoder.eval() 
        with torch.no_grad(): 
            ov_model = ov.convert_model(vae_decoder, example_input=latents, input=([1, 4, width, height],)) 
        ov.save_model(ov_model, ir_path) 
        del ov_model 
        cleanup_torchscript_cache() 
        print("VAE decoder successfully converted to IR") 

if not VAE_ENCODER_OV_PATH.exists(): 
    convert_vae_encoder(vae, VAE_ENCODER_OV_PATH, 768, 768) 
else: 
    print(f"VAE encoder will be loaded from {VAE_ENCODER_OV_PATH}") 

VAE_DECODER_OV_PATH = sd2_1_model_dir / "vae_decoder.xml" 

if not VAE_DECODER_OV_PATH.exists(): 
    convert_vae_decoder(vae, VAE_DECODER_OV_PATH, 96, 96) 
else: 
    print(f"VAE decoder will be loaded from {VAE_DECODER_OV_PATH}") 

del vae 
gc.collect();

import inspect 
from typing import List, Optional, Union, Dict 

import PIL 
import cv2 
import torch 

from transformers import CLIPTokenizer 
from diffusers import DiffusionPipeline 
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler 
    """ 
    Preprocessing helper function for calculating image size for resize with peserving original aspect ratio 
    and fitting image to specific window size 

    Parameters: 
        dst_width (int): destination window width 
        dst_height (int): destination window height 
        image_width (int): source image width 
        image_height (int): source image height 
    Returns: 
        result_width (int): calculated width for resize 
        result_height (int): calculated height for resize 
    """ 
    im_scale = min(dst_height / image_height, dst_width / image_width) 
    return int(im_scale * image_width), int(im_scale * image_height) 

def preprocess(image: PIL.Image.Image): 
    """ 
    Image preprocessing function.Takes image in PIL.Image format, resizes it to keep aspect ration and fits to model input window 512x512, 
    then converts it to np.ndarray and adds padding with zeros on right or bottom side of image (depends from aspect ratio), after that 
    converts data to float32 data type and change range of values from [0, 255] to [-1, 1], finally, converts data layout from planar NHWC to NCHW.The function returns preprocessed input tensor and padding size, which can be used in postprocessing.

    Parameters: 
        image (PIL.Image.Image): input image 
    Returns: 
        image (np.ndarray): preprocessed image tensor 
        meta (Dict): dictionary with preprocessing metadata info 
    """ 
    src_width, src_height = image.size 
    dst_width, dst_height = scale_fit_to_window(512, 512, src_width, src_height) 
    image = np.array(image.resize((dst_width, dst_height), resample=PIL.Image.Resampling.LANCZOS))[None, :] 
    pad_width = 512 - dst_width 
    pad_height = 512 - dst_height 
    pad = ((0, 0), (0, pad_height), (0, pad_width), (0, 0)) 
    image = np.pad(image, pad, mode="constant") 
    image = image.astype(np.float32) / 255.0 
    image = 2.0 * image - 1.0 image = image.transpose(0, 3, 1, 2) 
    return image, {"padding": pad, "src_width": src_width, "src_height": src_height} 

class OVStableDiffusionPipeline(DiffusionPipeline): 
    def __init__( 
        self, vae_decoder: ov.Model, 
        text_encoder: ov.Model, 
        tokenizer: CLIPTokenizer, 
        unet: ov.Model, 
        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], 
        vae_encoder: ov.Model = None, 
    ): 
        """ 
        Pipeline for text-to-image generation using Stable Diffusion.
        Parameters: 
            vae_decoder (Model):
                Variational Auto-Encoder (VAE) Model to decode images to and from latent representations. 
            text_encoder (Model): 
                Frozen text-encoder.Stable Diffusion uses the text portion of 
                [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically 
                the clip-vit-large-patch14(https://huggingface.co/openai/clip-vit-large-patch14) variant. 
            tokenizer (CLIPTokenizer):
                Tokenizer of class CLIPTokenizer(https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). 
            unet (Model):
                 Conditional U-Net architecture to denoise the encoded image latents. 
            vae_encoder (Model):
                Variational Auto-Encoder (VAE) Model to encode images to latent representation. 
            scheduler (SchedulerMixin): 
                A scheduler to be used in combination with unet to denoise the encoded image latents.Can be one of 
                DDIMScheduler, LMSDiscreteScheduler, or PNDMScheduler.
        """ 
        super().__init__() 
        self.scheduler = scheduler 
        self.vae_decoder = vae_decoder 
        self.vae_encoder = vae_encoder 
        self.text_encoder = text_encoder 
        self.unet = unet 
        self.register_to_config(unet=unet) 
        self._text_encoder_output = text_encoder.output(0) 
        self._unet_output = unet.output(0) 
        self._vae_d_output = vae_decoder.output(0) 
        self._vae_e_output = vae_encoder.output(0) if vae_encoder is not None else None 
        self.height = self.unet.input(0).shape[2] * 8 
        self.width = self.unet.input(0).shape[3] * 8 
        self.tokenizer = tokenizer 

    def __call__( 
        self, prompt: Union[str, List[str]], 
        image: PIL.Image.Image = None, 
        negative_prompt: Union[str, List[str]] = None, 
        num_inference_steps: Optional[int] = 50, 
        guidance_scale: Optional[float] = 7.5, 
        eta: Optional[float] = 0.0, 
        output_type: Optional[str] = "pil", 
        seed: Optional[int] = None, 
        strength: float = 1.0, 
    ): 
        """ 
        Function invoked when calling the pipeline for generation.
        Parameters: 
            prompt (str or List[str]):
                The prompt or prompts to guide the image generation. 
            image (PIL.Image.Image, *optional*, None):
                Intinal image for generation. 
            negative_prompt (str or List[str]):
                The negative prompt or prompts to guide the image generation. 
            num_inference_steps (int, *optional*, defaults to 50): 
                The number of denoising steps.More denoising steps usually lead to a higher quality image at the 
                expense of slower inference. 
            guidance_scale (float, *optional*, defaults to 7.5):
                Guidance scale as defined in Classifier-Free Diffusion Guidance(https://arxiv.org/abs/2207.12598). 
                guidance_scale is defined as `w` of equation 2.
                Higher guidance scale encourages to generate images that are closely linked to the text prompt, 
                usually at the expense of lower image quality. 
            eta (float, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502.Only applies to 
                [DDIMScheduler], will be ignored for others. 
            output_type (`str`, *optional*, defaults to "pil"): The output format of the generate image.Choose between 
            [PIL](https://pillow.readthedocs.io/en/stable/):
                PIL.Image.Image or np.array. 
            seed (int, *optional*, None):
                Seed for random generator state initialization. 
            strength (int, *optional*, 1.0): 
                strength between initial image and generated in Image-to-Image pipeline, do not used in Text-to-Image 
        Returns:
            Dictionary with keys: 
                sample - the last generated image PIL.Image.Image or np.array 
        """ 
        if seed is not None: 
            np.random.seed(seed) 
        # ここで、`guidance_scale` は、Imagen 論文の式 (2) のガイダンス重み `w` と 
        # 同様に定義されます: https://arxiv.org/pdf/2205.11487.pdf`guidance_scale = 1` 
        # 分類子のフリーガイダンスを行わないことに相当します 
        do_classifier_free_guidance = guidance_scale > 1.0 
        # プロンプトテキスト埋め込みを取得 
        text_embeddings = self._encode_prompt( 
            prompt, 
            do_classifier_free_guidance=do_classifier_free_guidance, 
            negative_prompt=negative_prompt, 
        ) 
        # タイムステップを設定 
        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()) 
        extra_set_kwargs = {} 
        if accepts_offset: 
            extra_set_kwargs["offset"] = 1 

        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs) 
        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength) 
        latent_timestep = timesteps[:1] 

        # ユーザーが指定しない限り、初期のランダムノイズを取得 
        latents, meta = self.prepare_latents(image, latent_timestep) 

        # すべてのスケジューラーが同じシグネチャーを持つわけではないので、スケジューラー・ステップに追加の kwargs を準備 
        # eta (η) は DDIMScheduler でのみ使用され、他のスケジューラーでは無視されます
        # eta は DDIM 論文の η に対応します: https://arxiv.org/abs/2010.02502 [0, 1] 
        # の範囲にある必要があります 
        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) 
        extra_step_kwargs = {} 
        if accepts_eta: 
            extra_step_kwargs["eta"] = eta 

        for t in self.progress_bar(timesteps):
            # 分類器のフリーガイダンスを行う場合は潜在変数を拡張 
            latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents 
            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) 

            # ノイズ残留を予測 
            noise_pred = self.unet([latent_model_input, np.array(t, dtype=np.float32), text_embeddings])[self._unet_output] 
            # ガイダンスを実行 
            if do_classifier_free_guidance: 
                noise_pred_uncond, noise_pred_text = noise_pred[0], noise_pred[1] 
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) 

            # 前のノイズサンプルを計算 x_t -> x_t-1 
            latents = self.scheduler.step(torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs)["prev_sample"].numpy() 
        # 画像の潜在変数を vae でスケールしてデコード 
        image = self.vae_decoder(latents * (1 / 0.18215))[self._vae_d_output] 

        image = self.postprocess_image(image, meta, output_type) 
        return {"sample": image} 

    def _encode_prompt( 
        self, 
        prompt: Union[str, List[str]], 
        num_images_per_prompt: int = 1, 
        do_classifier_free_guidance: bool = True, 
        negative_prompt: Union[str, List[str]] = None, 
    ): 
        """ 
        Encodes the prompt into text encoder hidden states.

        Parameters: 
            prompt (str or list(str)): prompt to be encoded 
            num_images_per_prompt (int): number of images that should be generated per prompt 
            do_classifier_free_guidance (bool): whether to use classifier free guidance or not 
            negative_prompt (str or list(str)): negative prompt to be encoded 
        Returns: 
            text_embeddings (np.ndarray): text encoder hidden states 
        """ 
        batch_size = len(prompt) if isinstance(prompt, list) else 1 

        # 入力プロンプトをトークン化 
        text_inputs = self.tokenizer( 
            prompt, 
            padding="max_length", 
            max_length=self.tokenizer.model_max_length, 
            truncation=True, 
            return_tensors="np", 
        ) 
        text_input_ids = text_inputs.input_ids 

        text_embeddings = self.text_encoder(text_input_ids)[self._text_encoder_output] 

        # プロンプトごとに各世代のテキスト埋め込みを複製 
        if num_images_per_prompt != 1: 
            bs_embed, seq_len, _ = text_embeddings.shape 
            text_embeddings = np.tile(text_embeddings, (1, num_images_per_prompt, 1)) 
            text_embeddings = np.reshape(text_embeddings, (bs_embed * num_images_per_prompt, seq_len, -1)) 

        # 分類器の無条件埋め込みを取得するフリーガイダンス 
        if do_classifier_free_guidance: 
            uncond_tokens: List[str] 
            max_length = text_input_ids.shape[-1] 
            if negative_prompt is None: 
                uncond_tokens = [""] * batch_size 
            elif isinstance(negative_prompt, str): 
                uncond_tokens = [negative_prompt] 
            else: 
                uncond_tokens = negative_prompt 
            uncond_input = self.tokenizer( 
                uncond_tokens, 
                padding="max_length", 
                max_length=max_length, 
                truncation=True, 
                return_tensors="np", 
            ) 

            uncond_embeddings = self.text_encoder(uncond_input.input_ids)[self.text_encoder_output] 

            # mps フレンドリーな方法を使用して、プロンプトごとに各世代の無条件埋め込みを複製 
            seq_len = uncond_embeddings.shape[1] 
            uncond_embeddings = np.tile(uncond_embeddings, (1, num_images_per_prompt, 1)) 
            uncond_embeddings = np.reshape(uncond_embeddings, (batch_size * num_images_per_prompt, seq_len, -1)) 

            # 分類器フリーのガイダンスでは、2 回のフォワードパスを実行する必要がある
            # ここでは、無条件埋め込みとテキスト埋め込みを 1 つのバッチに連結して、 
            # 2 回のフォワードパスを回避します 
            text_embeddings = np.concatenate([uncond_embeddings, text_embeddings]) 

        return text_embeddings 

    def prepare_latents(self, image: PIL.Image.Image = None, latent_timestep: torch.Tensor = None): 
        """ 
        Function for getting initial latents for starting generation 

        Parameters: 
            image (PIL.Image.Image, *optional*, None): Input image for generation, if not provided randon noise will be used as starting point 
            latent_timestep (torch.Tensor, *optional*, None): Predicted by scheduler initial step for image generation, required for latent image mixing with nosie 
        Returns: 
            latents (np.ndarray): Image encoded in latent space 
        """ 
        latents_shape = (1, 4, self.height // 8, self.width // 8) 
        noise = np.random.randn(*latents_shape).astype(np.float32) 
        if image is None:
            # LMSDiscreteScheduler を使用する場合は、潜在変数がシグマで乗算されていることを確認 
            if isinstance(self.scheduler, LMSDiscreteScheduler): 
                noise = noise * self.scheduler.sigmas[0].numpy() 
            return noise, {} 
        input_image, meta = preprocess(image) 
        latents = self.vae_encoder(input_image)[self._vae_e_output] 
        latents = latents * 0.18215 
        latents = self.scheduler.add_noise(torch.from_numpy(latents), torch.from_numpy(noise), latent_timestep).numpy() 
        return latents, meta 

    def postprocess_image(self, image: np.ndarray, meta: Dict, output_type: str = "pil"): 
        """ 
        Postprocessing for decoded image.Takes generated image decoded by VAE decoder, unpad it to initila image size (if required), 
        normalize and convert to [0, 255] pixels range.Optionally, convertes it from np.ndarray to PIL.Image format 

        Parameters: 
            image (np.ndarray): Generated image 
            meta (Dict): Metadata obtained on latents preparing step, can be empty 
            output_type (str, *optional*, pil): Output format for result, can be pil or numpy 
        Returns: 
            image (List of np.ndarray or PIL.Image.Image): Postprocessed images 
        """ 
        if "padding" in meta: 
            pad = meta["padding"] 
            (_, end_h), (_, end_w) = pad[1:3] 
            h, w = image.shape[2:] 
            unpad_h = h - end_h 
            unpad_w = w - end_w 
            image = image[:, :, :unpad_h, :unpad_w] 
        image = np.clip(image / 2 + 0.5, 0, 1) 
        image = np.transpose(image, (0, 2, 3, 1)) 
        # 9. PIL へ変換 
        if output_type == "pil": 
            image = self.numpy_to_pil(image) 
            if "src_height" in meta: 
                orig_height, orig_width = meta["src_height"], meta["src_width"] 
                image = [img.resize((orig_width, orig_height), PIL.Image.Resampling.LANCZOS) for img in image] 
        else: 
            if "src_height" in meta: 
                orig_height, orig_width = meta["src_height"], meta["src_width"] 
                image = [cv2.resize(img, (orig_width, orig_width)) for img in image] 
        return image 

    def get_timesteps(self, num_inference_steps: int, strength: float): 
        """ 
        Helper function for getting scheduler timesteps for generation 
        In case of image-to-image generation, it updates number of steps according to strength 

        Parameters: 
            num_inference_steps (int): 
                number of inference steps for generation 
            strength (float): 
                value between 0.0 and 1.0, that controls the amount of noise that is added to the input image.Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.
        """ 
        # init_timestep を使用して元のタイムステップを取得 
        init_timestep = min(int(num_inference_steps * strength), num_inference_steps) 

        t_start = max(num_inference_steps - init_timestep, 0) 
        timesteps = self.scheduler.timesteps[t_start:] 

        return timesteps, num_inference_steps - t_start

import ipywidgets as widgets 

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

device

ov_config = {"INFERENCE_PRECISION_HINT": "f32"} if device.value != "CPU" else {} 

text_enc = core.compile_model(TEXT_ENCODER_OV_PATH, device.value) 
unet_model = core.compile_model(UNET_OV_PATH, device.value) 
vae_decoder = core.compile_model(VAE_DECODER_OV_PATH, device.value, ov_config) 
vae_encoder = core.compile_model(VAE_ENCODER_OV_PATH, device.value, ov_config)

from transformers import CLIPTokenizer 

scheduler = DDIMScheduler.from_config(conf) # DDIMScheduler is used because UNet quantization produces better results with it 
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") 

ov_pipe = OVStableDiffusionPipeline( 
    tokenizer=tokenizer, 
    text_encoder=text_enc, 
    unet=unet_model, 
    vae_encoder=vae_encoder, 
    vae_decoder=vae_decoder, 
    scheduler=scheduler, 
)

%%skip not $to_quantize.value 

import datasets 
import numpy as np 
from tqdm.notebook import tqdm 
from typing import Any, Dict, List 

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 

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

    def __call__(self, *args, **kwargs): 
        if np.random.rand() <= self.keep_prob: 
            self.data_cache.append(*args) 
        return super().__call__(*args, **kwargs) 

def collect_calibration_data(ov_pipe, calibration_dataset_size: int, num_inference_steps: int) -> List[Dict]: 
    original_unet = ov_pipe.unet 
    calibration_data = [] 
    ov_pipe.unet = CompiledModelDecorator(original_unet, calibration_data, keep_prob=0.7) 
    disable_progress_bar(ov_pipe) 

    dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True).shuffle(seed=42) 

    # データ収集のため推論を実行 
    pbar = tqdm(total=calibration_dataset_size) 
    for batch in dataset: 
        prompt = batch["caption"] 
        if len(prompt) > ov_pipe.tokenizer.model_max_length: 
            continue 
        ov_pipe(prompt, num_inference_steps=num_inference_steps, seed=1) 
        pbar.update(len(calibration_data) - pbar.n) 
        if pbar.n >= calibration_dataset_size: 
            break 

    disable_progress_bar(ov_pipe, disable=False) 
    ov_pipe.unet = original_unet 
    return calibration_data

%%skip not $to_quantize.value 

from collections import deque 
from transformers import set_seed 
import nncf 

def get_operation_const_op(operation, const_port_id: int): 
    node = operation.input_value(const_port_id).get_node() 
    queue = deque([node]) 
    constant_node = None 
    allowed_propagation_types_list = ["Convert", "FakeQuantize", "Reshape"] 

    while len(queue) != 0: 
        curr_node = queue.popleft() 
        if curr_node.get_type_name() == "Constant": 
            constant_node = curr_node 
            break 
        if len(curr_node.inputs()) == 0: 
            break 
        if curr_node.get_type_name() in allowed_propagation_types_list: 
            queue.append(curr_node.input_value(0).get_node()) 

    return constant_node 

def is_embedding(node) -> bool: 
    allowed_types_list = ["f16", "f32", "f64"] 
    const_port_id = 0 
    input_tensor = node.input_value(const_port_id) 
    if input_tensor.get_element_type().get_type_name() in allowed_types_list: 
        const_node = get_operation_const_op(node, const_port_id) 
        if const_node is not None: 
            return True 

        return False 

def collect_ops_with_weights(model): 
    ops_with_weights = [] 
    for op in model.get_ops(): 
        if op.get_type_name() == "MatMul": 
            constant_node_0 = get_operation_const_op(op, const_port_id=0) 
            constant_node_1 = get_operation_const_op(op, const_port_id=1) 
            if constant_node_0 or constant_node_1: 
                ops_with_weights.append(op.get_friendly_name()) 
        if op.get_type_name() == "Gather" and is_embedding(op): 
            ops_with_weights.append(op.get_friendly_name()) 

    return ops_with_weights 

UNET_INT8_OV_PATH = sd2_1_model_dir / 'unet_optimized.xml' 
if not UNET_INT8_OV_PATH.exists(): 
    calibration_dataset_size = 300 
    set_seed(1) 
    unet_calibration_data = collect_calibration_data(ov_pipe, calibration_dataset_size=calibration_dataset_size, num_inference_steps=50) 

    unet = core.read_model(UNET_OV_PATH) 

    # 重みが圧縮される操作を収集 
    unet_ignored_scope = collect_ops_with_weights(unet) 

    # モデルの重みを圧縮 
    compressed_unet = nncf.compress_weights(unet, ignored_scope=nncf.IgnoredScope(types=['Convolution'])) 

    # 畳み込み層の重みと活性化の両方を量子化 
    quantized_unet = nncf.quantize( 
        model=compressed_unet, 
        calibration_dataset=nncf.Dataset(unet_calibration_data), 
        subset_size=calibration_dataset_size, 
        model_type=nncf.ModelType.TRANSFORMER, 
        ignored_scope=nncf.IgnoredScope(names=unet_ignored_scope), 

advanced_parameters=nncf.AdvancedQuantizationParameters(smooth_quant_alpha=-1) 
) 
ov.save_model(quantized_unet, UNET_INT8_OV_PATH)

%%skip not $to_quantize.value 

int8_unet_model = core.compile_model(UNET_INT8_OV_PATH, device.value) 
int8_ov_pipe = OVStableDiffusionPipeline( 
    tokenizer=tokenizer, 
    text_encoder=text_enc, 
    unet=int8_unet_model, 
    vae_encoder=vae_encoder, 
    vae_decoder=vae_decoder, 
    scheduler=scheduler 
)

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

%%skip not $to_quantize.value 

import time 

def calculate_inference_time(pipeline, validation_data): 
    inference_time = [] 
    pipeline.set_progress_bar_config(disable=True) 
    for prompt in validation_data: 
        start = time.perf_counter() 
        _ = pipeline(prompt, num_inference_steps=10, seed=0) 
        end = time.perf_counter() 
        delta = end - start 
        inference_time.append(delta) 
    return np.median(inference_time)

%%skip not $to_quantize.value 

validation_size = 10 
validation_dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", streaming=True, trust_remote_code=True).take(validation_size) 
validation_data = [batch["caption"] for batch in validation_dataset] 

fp_latency = calculate_inference_time(ov_pipe, validation_data) 
int8_latency = calculate_inference_time(int8_ov_pipe, validation_data) 
print(f"Performance speed-up: {fp_latency / int8_latency:.3f}")