LCM LoRA および ControlNet コンディショニングによるテキストからの画像#

%pip install -q "torch" transformers "diffusers>=0.22.0" "controlnet-aux>=0.0.6" "peft==0.6.2" accelerate --extra-index-url https://download.pytorch.org/whl/cpu 
%pip install -q "openvino>=2023.2.0" pillow "gradio>=4.19" "datasets>=2.14.6" "nncf>=2.7.0"

from pathlib import Path 

controlnet_id = "lllyasviel/control_v11p_sd15_normalbae" 
adapter_id = "latent-consistency/lcm-lora-sdv1-5" 
stable_diffusion_id = "runwayml/stable-diffusion-v1-5" 

TEXT_ENCODER_OV_PATH = Path("model/text_encoder.xml") 
UNET_OV_PATH = Path("model/unet_controlnet.xml") 
CONTROLNET_OV_PATH = Path("model/controlnet-normalbae.xml") 
VAE_DECODER_OV_PATH = Path("model/vae_decoder.xml") 
TOKENIZER_PATH = Path("model/tokenizer") 
SCHEDULER_PATH = Path("model/scheduler") 

skip_models = TEXT_ENCODER_OV_PATH.exists() and UNET_OV_PATH.exists() and CONTROLNET_OV_PATH.exists() and VAE_DECODER_OV_PATH.exists()

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel 
import gc 

def load_original_pytorch_pipeline_components(controlnet_id: str, stable_diffusion_id: str, adapter_id: str): 
    """ 
    Helper function for loading Stable Diffusion ControlNet pipeline and applying LCM LoRA 

    Parameters: 
        controlnet_id: model id from HuggingFace hub or local path for loading ControlNet model 
        stable_diffusion_id: model id from HuggingFace hub or local path for loading Stable Diffusion model 
        adapter_id: LCM LoRA id from HuggingFace hub or local path 
    Returns: 
        controlnet: ControlNet model 
        text_encoder: Stable Diffusion Text Encoder 
        unet: Stable Diffusion U-Net 
        vae: Stable Diffusion Variational Autoencoder (VAE) 
    """ 

    # 制御ネットモデルをロード 
    controlnet = ControlNetModel.from_pretrained(controlnet_id) 
    # stable diffusion パイプラインをロード 
    pipe = StableDiffusionControlNetPipeline.from_pretrained(stable_diffusion_id, controlnet=controlnet) 
    # LCM LoRA の重みをロード 
    pipe.load_lora_weights(adapter_id) 
    # LoRA の重みを UNet と融合 
    pipe.fuse_lora() text_encoder = pipe.text_encoder 
    text_encoder.eval() 
    unet = pipe.unet 
    unet.eval() 
    vae = pipe.vae 
    vae.eval() 
    del pipe 
    gc.collect() 
    return controlnet, text_encoder, unet, vae

controlnet, text_encoder, unet, vae = None, None, None, None 
if not skip_models: 
    controlnet, text_encoder, unet, vae = load_original_pytorch_pipeline_components(controlnet_id, stable_diffusion_id, adapter_id)

from controlnet_aux import NormalBaeDetector 
from diffusers.utils import load_image 
import requests 
import matplotlib.pyplot as plt 
from PIL import Image 
import numpy as np 

example_image_url = 
"https://huggingface.co/lllyasviel/control_v11p_sd15_normalbae/resolve/main/images/input.png" 
r = requests.get(example_image_url) 
with open("example.png", "wb") as f: 
    f.write(r.content) 

processor = NormalBaeDetector.from_pretrained("lllyasviel/Annotators") 

image = load_image("example.png") 
control_image = processor(image) 

def visualize_results( 
    orig_img: Image.Image, 
    normal_img: Image.Image, 
    result_img: Image.Image = None, 
    save_fig: bool = False, 
): 
    """ 
    Helper function for results visualization 

    Parameters: 
        orig_img (Image.Image): original image 
        normal_img (Image.Image): image with bwith surface normal information 
        result_img (Image.Image, optional, default None): generated image 
        safe_fig (bool, optional, default False): allow saving visualization result on disk 
    Returns: 
        fig (matplotlib.pyplot.Figure): matplotlib generated figure contains drawing result 
    """ 
    orig_title = "Original image" 
    control_title = "Normal map" 
    orig_img = orig_img.resize(normal_img.size if result_img is None else result_img.size) 
    im_w, im_h = orig_img.size 
    is_horizontal = im_h <= im_w 
    figsize = (20, 20) 
    num_images = 3 if result_img is not None else 2 
    fig, axs = plt.subplots( 
        num_images if is_horizontal else 1, 
        1 if is_horizontal else num_images, 
        figsize=figsize, 
        sharex="all", 
        sharey="all", 
    ) 
    fig.patch.set_facecolor("white") 
    list_axes = list(axs.flat) 
    for a in list_axes: 
        a.set_xticklabels([]) 
        a.set_yticklabels([]) 
        a.get_xaxis().set_visible(False) 
        a.get_yaxis().set_visible(False) 
        a.grid(False) 
    list_axes[0].imshow(np.array(orig_img)) 
    list_axes[1].imshow(np.array(normal_img)) 
    list_axes[0].set_title(orig_title, fontsize=15) 
    list_axes[1].set_title(control_title, fontsize=15) 
    if result_img is not None: 
        list_axes[2].imshow(np.array(result_img)) 
        list_axes[2].set_title("Result", fontsize=15) 

    fig.subplots_adjust(wspace=0.01 if is_horizontal else 0.00, hspace=0.01 if is_horizontal else 0.1) 
    fig.tight_layout() 
    if save_fig: 
        fig.savefig("result.png", bbox_inches="tight") 
    return fig 

fig = visualize_results(image, control_image)

from diffusers import DiffusionPipeline 
from transformers import CLIPTokenizer 
from typing import Union, List, Optional, Tuple 
import cv2 

def scale_fit_to_window(dst_width: int, dst_height: int, image_width: int, image_height: int): 
    """ 
    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: Image.Image, dst_height: int = 512, dst_width: int = 512): 
    """ 
    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 (Image.Image): input image 
        dst_width: destination image width 
        dst_height: destination image height 
    Returns: 
        image (np.ndarray): preprocessed image tensor 
        pad (Tuple[int]): pading size for each dimension for restoring image size in postprocessing 
    """ 
    src_width, src_height = image.size 
    res_width, res_height = scale_fit_to_window(dst_width, dst_height, src_width, src_height) 
    image = np.array(image.resize((res_width, res_height), resample=Image.Resampling.LANCZOS))[None, :] 
    pad_width = dst_width - res_width 
    pad_height = dst_height - res_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 = image.transpose(0, 3, 1, 2) 
    return image, pad 

def randn_tensor( 
    shape: Union[Tuple, List], 
    dtype: Optional[torch.dtype] = torch.float32, 
): 
    """ 
    Helper function for generation random values tensor with given shape and data type 

    Parameters: 
        shape (Union[Tuple, List]): shape for filling random values 
        dtype (torch.dtype, *optiona*, torch.float32): data type for result 
    Returns: 
        latents (np.ndarray): tensor with random values with given data type and shape (usually represents noise in latent space) 
    """ 
    latents = torch.randn(shape, dtype=dtype) 
    return latents.numpy() 

class OVControlNetStableDiffusionPipeline(DiffusionPipeline): 
    """ 
    OpenVINO inference pipeline for Stable Diffusion with ControlNet guidence 
    """ 

    def __init__( 
        self, 
        tokenizer: CLIPTokenizer, 
        scheduler, 
        core: ov.Core, 
        controlnet: ov.Model, 
        text_encoder: ov.Model, 
        unet: ov.Model, 
        vae_decoder: ov.Model, 
        device: str = "AUTO", 
    ): 
        super().__init__() 
        self.tokenizer = tokenizer 
        self.vae_scale_factor = 8 
        self.scheduler = scheduler 
        self.load_models(core, device, controlnet, text_encoder, unet, vae_decoder) 

    def load_models( 
        self, 
        core: ov.Core, 
        device: str, 
        controlnet: ov.Model, 
        text_encoder: ov.Model, 
        unet: ov.Model, 
        vae_decoder: ov.Model, 
    ): 
        """ 
        Function for loading models on device using OpenVINO 

        Parameters: 
            core (Core): OpenVINO runtime Core class instance 
            device (str): inference device 
            controlnet (Model): OpenVINO Model object represents ControlNet 
            text_encoder (Model): OpenVINO Model object represents text encoder 
            unet (Model): OpenVINO Model object represents UNet 
            vae_decoder (Model): OpenVINO Model object represents vae decoder 
        Returns 
            None 
        """ 
        self.text_encoder = core.compile_model(text_encoder, device) 
        self.register_to_config(controlnet=core.compile_model(controlnet, device)) 
        self.register_to_config(unet=core.compile_model(unet, device)) 
        ov_config = {"INFERENCE_PRECISION_HINT": "f32"} if device != "CPU" else {} 
        self.vae_decoder = core.compile_model(vae_decoder, device, ov_config) 

    def __call__( 
        self, 
        prompt: Union[str, List[str]], 
        image: Image.Image, 
        num_inference_steps: int = 4, 
        height: int = 512, 
        width: int = 512, 
        negative_prompt: Union[str, List[str]] = None, 
        guidance_scale: float = 0.5, 
        controlnet_conditioning_scale: float = 1.0, 
        latents: Optional[np.array] = None, 
        output_type: Optional[str] = "pil", 
    ): 
        """ 
        Function invoked when calling the pipeline for generation.

        Parameters: 
            prompt (`str` or `List[str]`):
                The prompt or prompts to guide the image generation. 
            image (`Image.Image`):
                `Image`, or tensor representing an image batch which will be repainted according to `prompt`. 
            num_inference_steps (`int`, *optional*, defaults to 100): 
                The number of denoising steps.More denoising steps usually lead to a higher quality image at the 
                expense of slower inference. 
            height (int, *optional*, defaults to 512): generated image height 
            width (int, *optional*, defaults to 512): generated image width 
            negative_prompt (`str` or `List[str]`): 
                negative prompt or prompts for generation 
            guidance_scale (`float`, *optional*, defaults to 0.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. of [Imagen 
                Paper](https://arxiv.org/pdf/2205.11487.pdf).Guidance scale is enabled by setting `guidance_scale > 
                1`.Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, 
                usually at the expense of lower image quality.This pipeline requires a value of at least `1`. 
            latents (`np.ndarray`, *optional*):
                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image 
                generation.Can be used to tweak the same generation with different prompts.If not provided, a latents 
                tensor will ge generated by sampling using the supplied random `generator`. 
            output_type (`str`, *optional*, defaults to `"pil"`): 
                The output format of the generate image.Choose between 
                [PIL](https://pillow.readthedocs.io/en/stable/): `Image.Image` or `np.array`.
        Returns: 
            image ([List[Union[np.ndarray, Image.Image]]): generaited images 
        """ 

        # 1. 呼び出しパラメーターを定義 
        batch_size = 1 if isinstance(prompt, str) else len(prompt)
        if guidance_scale < 1 and negative_prompt: 
            guidance_scale += 1
        # ここで、`guidance_scale` は、Imagen 論文の式 (2) のガイダンス重み `w` と同様に定義されます: 
        # https://arxiv.org/pdf/2205.11487.pdf`guidance_scale = 1` 
        # は、分類器フリーのガイダンスを行わないことに相当 
        do_classifier_free_guidance = guidance_scale > 1.0 
        # 2. 入力プロンプトをエンコード 
        text_embeddings = self._encode_prompt( 
            prompt, 
            do_classifier_free_guidance=do_classifier_free_guidance, 
            negative_prompt=negative_prompt, 
        ) 

        # 3. 画像の前処理 
        orig_width, orig_height = image.size 
        image, pad = preprocess(image, height, width) 
        if do_classifier_free_guidance: 
            image = np.concatenate(([image] * 2)) 

        # 4. タイムステップを設定 
        self.scheduler.set_timesteps(num_inference_steps) 
        timesteps = self.scheduler.timesteps 

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

        # 6 .ノイズ除去ループ 
        with self.progress_bar(total=num_inference_steps) as progress_bar: 
            for i, t in enumerate(timesteps): 
                # 分類器フリーのガイダンスを実行する場合は潜在変数を拡張
                # pix2pix ではテキストと入力画像の両方にガイダンスが適用されるため、 
                # 潜在変数は 3 倍に拡張されます 
                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) 

                result = self.controlnet( 
                    [latent_model_input, t, text_embeddings, image], 
                    share_inputs=True, 
                    share_outputs=True, 
                ) 
                down_and_mid_blok_samples = [sample * controlnet_conditioning_scale for _, sample in result.items()] 

                # ノイズ残留を予測 
                noise_pred = self.unet( 
                    [ 
                        latent_model_input, 
                        t, 
                        text_embeddings, 
                        *down_and_mid_blok_samples, 
                    ], 
                    share_inputs=True, 
                    share_outputs=True, 
                )[0] 

                # ガイダンスを実行 
                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)).prev_sample.numpy() 
                progress_bar.update() 

        # 7. 後処理 
        image = self.decode_latents(latents, pad) 

        # 8. PIL に変換 
        if output_type == "pil": 
            image = self.numpy_to_pil(image) 
            image = [img.resize((orig_width, orig_height), Image.Resampling.LANCZOS) for img in image] 
        else: 
            image = [cv2.resize(img, (orig_width, orig_width)) for img in image] 

        return 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, share_inputs=True, share_outputs=True)[0] 

        # プロンプトごとに各世代のテキスト埋め込みを複製 
        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, share_inputs=True, share_outputs=True)[0] 

            # 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, 
        batch_size: int, 
        num_channels_latents: int, 
        height: int, 
        width: int, 
        dtype: np.dtype = torch.float32, 
        latents: np.ndarray = None, 
    ): 
        """ 
        Preparing noise to image generation.If initial latents are not provided, they will be generated randomly, 
        then prepared latents scaled by the standard deviation required by the scheduler 

        Parameters: 
            batch_size (int): input batch size 
            num_channels_latents (int): number of channels for noise generation 
            height (int): image height 
            width (int): image width 
            dtype (np.dtype, *optional*, np.float32): dtype for latents generation 
            latents (np.ndarray, *optional*, None): initial latent noise tensor, if not provided will be generated 
        Returns: 
            latents (np.ndarray): scaled initial noise for diffusion 
        """ 
        shape = ( 
            batch_size, 
            num_channels_latents, 
            height // self.vae_scale_factor, 
            width // self.vae_scale_factor, 
        ) 
        if latents is None: 
            latents = randn_tensor(shape, dtype=dtype) 
        else: 
            latents = latents 

        # スケジューラーが要求する標準偏差で初期ノイズをスケール 
        latents = latents * self.scheduler.init_noise_sigma 
        return latents 

    def decode_latents(self, latents: np.array, pad: Tuple[int]): 
        """ 
        Decode predicted image from latent space using VAE Decoder and unpad image result 

        Parameters: 
            latents (np.ndarray): image encoded in diffusion latent space 
            pad (Tuple[int]): each side padding sizes obtained on preprocessing step 
        Returns: 
            image: decoded by VAE decoder image 
        """ 
        latents = 1 / 0.18215 * latents 
        image = self.vae_decoder(latents)[0] 
        (_, 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)) 
        return image

prompt = "A head full of roses" 
torch.manual_seed(4257) 

result = ov_pipe(prompt, control_image, 4) 
result[0]

fig = visualize_results(image, control_image, result[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

# `skip_kernel_extension` モジュールを取得 
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) 

int8_pipe = None 

%load_ext skip_kernel_extension

import gradio as gr 

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

quantized_model_present = int8_pipe is not None 

gr.close_all() 
    with gr.Blocks() as demo: 
        with gr.Row(): 
            with gr.Column(): 
                inp_img = gr.Image(label="Input image") 
            with gr.Column(visible=True) as step1: 
                out_normal = gr.Image(label="Normal Map", type="pil", interactive=False) 
                btn = gr.Button() 
                inp_prompt = gr.Textbox(label="Prompt") 
                inp_neg_prompt = gr.Textbox( 
                    "", 
                    label="Negative prompt", 
                ) 
                with gr.Accordion("Advanced options", open=False): 
                    guidance_scale = gr.Slider( 
                        label="Guidance scale", 
                        minimum=0.1, 
                        maximum=2, 
                        step=0.1, 
                        value=0.5, 
                    ) 
                    inp_seed = gr.Slider(label="Seed", value=42, maximum=MAX_SEED) 
                    inp_steps = gr.Slider(label="Steps", value=4, minimum=1, maximum=50, step=1) 
        with gr.Column(visible=True) as step2: 
            out_result = gr.Image(label="Result (Original)") 
        with gr.Column(visible=quantized_model_present) as quantization_step: 
            int_result = gr.Image(label="Result (Quantized)") 
    examples = gr.Examples([["example.png", "a head full of roses"]], [inp_img, inp_prompt]) 

    def extract_normal_map(img): 
        if img is None: 
            raise gr.Error("Please upload the image or use one from the examples list") 
        return processor(img) 

    def generate(img, prompt, negative_prompt, seed, num_steps, guidance_scale): 
        torch.manual_seed(seed) 
        control_img = extract_normal_map(img) 

        result = ov_pipe( 
            prompt, 
            control_img, 
            num_steps, 
            guidance_scale=guidance_scale, 
            negative_prompt=negative_prompt, 
        )[0] 
        if int8_pipe is not None: 
            torch.manual_seed(seed) 
            int8_result = int8_pipe( 
                prompt, 
                control_img, 
                num_steps, 
                guidance_scale=guidance_scale, 
                negative_prompt=negative_prompt, 
            )[0] 
            return control_img, result, int8_result 
        return control_img, result 

    output_images = [out_normal, out_result] 
    if quantized_model_present: 
        output_images.append(int_result) 
    btn.click( 
        generate, 
        [inp_img, inp_prompt, inp_neg_prompt, inp_seed, inp_steps, guidance_scale], 
        output_images, 
    ) 

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