潜在整合性モデルと OpenVINO による画像生成#

%pip install -q "torch>=2.1" --index-url https://download.pytorch.org/whl/cpu 
%pip install -q "openvino>=2023.1.0" transformers "diffusers>=0.23.1" pillow "gradio>=4.19" "nncf>=2.7.0" "datasets>=2.14.6" "peft==0.6.2" --extra-index-url https://download.pytorch.org/whl/cpu

import gc 
import warnings 
from pathlib import Path 
from diffusers import DiffusionPipeline 
import numpy as np 

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=False): 
    pipe = DiffusionPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7") 
    scheduler = pipe.scheduler 
    tokenizer = pipe.tokenizer 
    feature_extractor = pipe.feature_extractor if not skip_safety_checker else None 
    safety_checker = pipe.safety_checker if not skip_safety_checker else None 
    text_encoder, unet, vae = None, None, None 
    if not skip_models: 
        text_encoder = pipe.text_encoder 
        text_encoder.eval() 
        unet = pipe.unet 
        unet.eval() 
        vae = pipe.vae 
        vae.eval() 
    del pipe 
    gc.collect() 
    return ( 
        scheduler, 
        tokenizer, 
        feature_extractor, 
        safety_checker, 
        text_encoder, 
        unet, 
        vae, 
    )

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)

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): 
    def __init__( 
        self, 
        vae_decoder: ov.Model, 
        text_encoder: ov.Model, 
        tokenizer: CLIPTokenizer, 
        unet: ov.Model, 
        scheduler: None, 
        safety_checker: StableDiffusionSafetyChecker, 
        feature_extractor: CLIPImageProcessor, 
        requires_safety_checker: bool = True, 
    ): 
        super().__init__() 
        self.vae_decoder = vae_decoder 
        self.text_encoder = text_encoder 
        self.tokenizer = tokenizer self.register_to_config(unet=unet) 
        self.scheduler = scheduler 
        self.safety_checker = safety_checker 
        self.feature_extractor = feature_extractor 
        self.vae_scale_factor = 2**3 
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) 

    def _encode_prompt( 
        self, 
        prompt, 
        num_images_per_prompt, 
        prompt_embeds: None, 
    ): 
        r""" 
        Encodes the prompt into text encoder hidden states.
        Args: 
            prompt (`str` or `List[str]`, *optional*): 
                prompt to be encoded 
            num_images_per_prompt (`int`): 
                number of images that should be generated per prompt 
            prompt_embeds (`torch.FloatTensor`, *optional*): 
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not 
                provided, text embeddings will be generated from `prompt` input argument.
        """ 

        if prompt_embeds is None: 
            text_inputs = self.tokenizer( 
                prompt, 
                padding="max_length", 
                max_length=self.tokenizer.model_max_length, 
                truncation=True, 
                return_tensors="pt", 
            ) 
            text_input_ids = text_inputs.input_ids 

            prompt_embeds = self.text_encoder(text_input_ids, share_inputs=True, share_outputs=True) 
            prompt_embeds = torch.from_numpy(prompt_embeds[0]) 

        bs_embed, seq_len, _ = prompt_embeds.shape 
        # プロンプトごとに各世代のテキスト埋め込みを複製 
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) 
        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) 

        # LCM ガイド蒸留により、無条件プロンプト埋め込みを取得する必要はありません 
        return prompt_embeds 

    def run_safety_checker(self, image, dtype): 
        if self.safety_checker is None: 
            has_nsfw_concept = None 
        else: 
            if torch.is_tensor(image): 
                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") 
            else: 
                feature_extractor_input = self.image_processor.numpy_to_pil(image) 
            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt") 
            image, has_nsfw_concept = self.safety_checker(images=image, clip_input=safety_checker_input.pixel_values.to(dtype)) 
        return image, has_nsfw_concept 

    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, latents=None): 
        shape = ( 
            batch_size, 
            num_channels_latents, 
            height // self.vae_scale_factor, 
            width // self.vae_scale_factor, 
        ) 
        if latents is None: 
            latents = torch.randn(shape, dtype=dtype) 
        # スケジューラーが要求する標準偏差で初期ノイズをスケール 
        latents = latents * self.scheduler.init_noise_sigma 
        return latents 

    def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32): 
        """ 
        see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 
        Args: 
        timesteps: torch.Tensor: generate embedding vectors at these timesteps 
        embedding_dim: int: dimension of the embeddings to generate 
        dtype: data type of the generated embeddings 
        Returns: 
        embedding vectors with shape `(len(timesteps), embedding_dim)` 
        """ 
        assert len(w.shape) == 1 
        w = w * 1000.0 

        half_dim = embedding_dim // 2 
        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) 
        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) 
        emb = w.to(dtype)[:, None] * emb[None, :] 
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) 
        if embedding_dim % 2 == 1: # ゼロパディング 
            emb = torch.nn.functional.pad(emb, (0, 1)) 
            assert emb.shape == (w.shape[0], embedding_dim) 
            return emb 

    @torch.no_grad() 
    def __call__( 
        self, 
        prompt: Union[str, List[str]] = None, 
        height: Optional[int] = 512, 
        width: Optional[int] = 512, 
        guidance_scale: float = 7.5, 
        num_images_per_prompt: Optional[int] = 1, 
        latents: Optional[torch.FloatTensor] = None, 
        num_inference_steps: int = 4, 
        lcm_origin_steps: int = 50, 
        prompt_embeds: Optional[torch.FloatTensor] = None, 
        output_type: Optional[str] = "pil", 
        return_dict: bool = True, 
        cross_attention_kwargs: Optional[Dict[str, Any]] = None, 
    ): 
        # 1. 呼び出しパラメーターを定義 
        if prompt is not None and isinstance(prompt, str): 
            batch_size = 1 
        elif prompt is not None and isinstance(prompt, list): 
            batch_size = len(prompt) 
        else: 
            batch_size = prompt_embeds.shape[0] 

        # do_classifier_free_guidance = guidance_scale > 0.0 
        # LCM 実装の場合: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG) 

        # 2. 入力プロンプトをエンコード 
        prompt_embeds = self._encode_prompt( 
            prompt, 
            num_images_per_prompt, 
            prompt_embeds=prompt_embeds, 
        ) 

        # 3. タイムステップを準備 
        self.scheduler.set_timesteps(num_inference_steps, original_inference_steps=lcm_origin_steps) 
        timesteps = self.scheduler.timesteps 

        # 4. 潜在変数を準備 
        num_channels_latents = 4 
        latents = self.prepare_latents( 
            batch_size * num_images_per_prompt, 
            num_channels_latents, 
            height, 
            width, 
            prompt_embeds.dtype, 
            latents, 
        ) 

        bs = batch_size * num_images_per_prompt 

        # 5. ガイダンススケールの埋め込みを取得 
        w = torch.tensor(guidance_scale).repeat(bs) 
        w_embedding = self.get_w_embedding(w, embedding_dim=256) 

        # 6. LCM マルチステップ・サンプリング・ループ: 
        with self.progress_bar(total=num_inference_steps) as progress_bar: 
            for i, t in enumerate(timesteps): 
                ts = torch.full((bs,), t, dtype=torch.long) 

                # モデル予測 (v-prediction, eps, x) 
                model_pred = self.unet( 
                    [latents, ts, prompt_embeds, w_embedding], 
                    share_inputs=True, 
                    share_outputs=True, 
                )[0] 

                # 前のノイズサンプルを計算 x_t -> x_t-1 
                latents, denoised = self.scheduler.step(torch.from_numpy(model_pred), t, latents, return_dict=False) 
                progress_bar.update() 

        if not output_type == "latent": 
            image = torch.from_numpy(self.vae_decoder(denoised / 0.18215, share_inputs=True, share_outputs=True)[0]) 
            image, has_nsfw_concept = self.run_safety_checker(image, prompt_embeds.dtype) 
        else: 
                image = denoised 
                has_nsfw_concept = None 

        if has_nsfw_concept is None: 
            do_denormalize = [True] * image.shape[0] 
        else: 
            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] 

        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) 

        if not return_dict: 
            return (image, has_nsfw_concept) 

        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

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

images[0]

skip_for_device = "GPU" in device.value 
to_quantize = widgets.Checkbox(value=not skip_for_device, description="Quantization", disabled=skip_for_device) 
to_quantize

int8_pipe = None 

# `skip_kernel_extension` モジュールを取得 
import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py", 
) 
open("skip_kernel_extension.py", "w").write(r.text) 
%load_ext skip_kernel_extension

import random 
import gradio as gr 
from functools import partial 

MAX_SEED = np.iinfo(np.int32).max 

examples = [ 
    "portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour," 
    "style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography", 
    "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k", 
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", 
    "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
 ] 

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: 
    if randomize_seed: 
        seed = random.randint(0, MAX_SEED) 
    return seed

MAX_IMAGE_SIZE = 768 

def generate( 
    pipeline: OVLatentConsistencyModelPipeline, 
    prompt: str, 
    seed: int = 0, 
    width: int = 512, 
    height: int = 512, 
    guidance_scale: float = 8.0, 
    num_inference_steps: int = 4, 
    randomize_seed: bool = False, 
    num_images: int = 1, 
    progress=gr.Progress(track_tqdm=True), 
): 
    seed = randomize_seed_fn(seed, randomize_seed) 
    torch.manual_seed(seed) 
    result = pipeline( 
        prompt=prompt, 
        width=width, 
        height=height, 
        guidance_scale=guidance_scale, 
        num_inference_steps=num_inference_steps, 
        num_images_per_prompt=num_images, 
        lcm_origin_steps=50, 
        output_type="pil", 
).images[0] 
return result, seed 

generate_original = partial(generate, ov_pipe) 
generate_optimized = partial(generate, int8_pipe) 
quantized_model_present = int8_pipe is not None 

with gr.Blocks() as demo: 
    with gr.Group(): 
        with gr.Row(): 
            prompt = gr.Text( 
                label="Prompt", 
                show_label=False, 
                max_lines=1, 
                placeholder="Enter your prompt", 
                container=False, 
            ) 
        with gr.Row(): 
            with gr.Column(): 
                result = gr.Image( 
                    label="Result (Original)" if quantized_model_present else "Image", 
                    type="pil", 
                ) 
                run_button = gr.Button("Run") 
            with gr.Column(visible=quantized_model_present): 
                result_optimized = gr.Image( 
                    label="Result (Optimized)", 
                    type="pil", 
                    visible=quantized_model_present, 
                ) 
                run_quantized_button = gr.Button(value="Run quantized", visible=quantized_model_present) 

    with gr.Accordion("Advanced options", open=False): 
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True) 
        randomize_seed = gr.Checkbox(label="Randomize seed across runs", value=True) 
        with gr.Row(): 
            width = gr.Slider( 
                label="Width", 
                minimum=256, 
                maximum=MAX_IMAGE_SIZE, 
                step=32, 
                value=512, 
            ) 
            height = gr.Slider( 
                label="Height", 
                minimum=256, 
                maximum=MAX_IMAGE_SIZE, 
                step=32, 
                value=512, 
            ) with gr.Row(): 
                guidance_scale = gr.Slider( 
                    label="Guidance scale for base", 
                    minimum=2, 
                    maximum=14, 
                    step=0.1, 
                    value=8.0, 
                ) 
                num_inference_steps = gr.Slider( 
                    label="Number of inference steps for base", 
                    minimum=1, 
                    maximum=8, 
                    step=1, 
                    value=4, 
                ) 

    gr.Examples( 
        examples=examples, 
        inputs=prompt, 
        outputs=result, 
        cache_examples=False, 
    ) 

    gr.on( 
        triggers=[ 
            prompt.submit, 
            run_button.click, 
        ], 
        fn=generate_original, 
        inputs=[ 
            prompt, 
            seed, 
            width, 
            height, 
            guidance_scale, 
            num_inference_steps, 
            randomize_seed, 
        ], 
        outputs=[result, seed], 
    ) 

    if quantized_model_present: 
        gr.on( 
            triggers=[ 
                prompt.submit, 
                run_quantized_button.click, 
            ], 
            fn=generate_optimized, 
            inputs=[ 
                prompt, 
                seed, 
                width, 
                height, 
                guidance_scale, 
                num_inference_steps, 
                randomize_seed, 
            ], 
            outputs=[result_optimized, seed], 
        )

try: 
    demo.queue().launch(debug=False) 
except Exception: 
    demo.queue().launch(share=True, debug=False) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/