diffusers/community-pipelines-mirror

1import argparse2import atexit3import inspect4import os5import time6import warnings7from typing import Any, Callable, Dict, List, Optional, Union8 9import numpy as np10import PIL.Image11import pycuda.driver as cuda12import tensorrt as trt13import torch14from PIL import Image15from pycuda.tools import make_default_context16from transformers import CLIPTokenizer17 18from diffusers import OnnxRuntimeModel, StableDiffusionImg2ImgPipeline, UniPCMultistepScheduler19from diffusers.image_processor import VaeImageProcessor20from diffusers.pipelines.pipeline_utils import DiffusionPipeline21from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput22from diffusers.schedulers import KarrasDiffusionSchedulers23from diffusers.utils import (24    deprecate,25    logging,26    replace_example_docstring,27)28from diffusers.utils.torch_utils import randn_tensor29 30 31# Initialize CUDA32cuda.init()33context = make_default_context()34device = context.get_device()35atexit.register(context.pop)36 37logger = logging.get_logger(__name__)  # pylint: disable=invalid-name38 39 40def load_engine(trt_runtime, engine_path):41    with open(engine_path, "rb") as f:42        engine_data = f.read()43    engine = trt_runtime.deserialize_cuda_engine(engine_data)44    return engine45 46 47class TensorRTModel:48    def __init__(49        self,50        trt_engine_path,51        **kwargs,52    ):53        cuda.init()54        stream = cuda.Stream()55        TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE)56        trt.init_libnvinfer_plugins(TRT_LOGGER, "")57        trt_runtime = trt.Runtime(TRT_LOGGER)58        engine = load_engine(trt_runtime, trt_engine_path)59        context = engine.create_execution_context()60 61        # allocates memory for network inputs/outputs on both CPU and GPU62        host_inputs = []63        cuda_inputs = []64        host_outputs = []65        cuda_outputs = []66        bindings = []67        input_names = []68        output_names = []69 70        for binding in engine:71            datatype = engine.get_binding_dtype(binding)72            if datatype == trt.DataType.HALF:73                dtype = np.float1674            else:75                dtype = np.float3276 77            shape = tuple(engine.get_binding_shape(binding))78            host_mem = cuda.pagelocked_empty(shape, dtype)79            cuda_mem = cuda.mem_alloc(host_mem.nbytes)80            bindings.append(int(cuda_mem))81 82            if engine.binding_is_input(binding):83                host_inputs.append(host_mem)84                cuda_inputs.append(cuda_mem)85                input_names.append(binding)86            else:87                host_outputs.append(host_mem)88                cuda_outputs.append(cuda_mem)89                output_names.append(binding)90 91        self.stream = stream92        self.context = context93        self.engine = engine94 95        self.host_inputs = host_inputs96        self.cuda_inputs = cuda_inputs97        self.host_outputs = host_outputs98        self.cuda_outputs = cuda_outputs99        self.bindings = bindings100        self.batch_size = engine.max_batch_size101 102        self.input_names = input_names103        self.output_names = output_names104 105    def __call__(self, **kwargs):106        context = self.context107        stream = self.stream108        bindings = self.bindings109 110        host_inputs = self.host_inputs111        cuda_inputs = self.cuda_inputs112        host_outputs = self.host_outputs113        cuda_outputs = self.cuda_outputs114 115        for idx, input_name in enumerate(self.input_names):116            _input = kwargs[input_name]117            np.copyto(host_inputs[idx], _input)118            # transfer input data to the GPU119            cuda.memcpy_htod_async(cuda_inputs[idx], host_inputs[idx], stream)120 121        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)122 123        result = {}124        for idx, output_name in enumerate(self.output_names):125            # transfer predictions back from the GPU126            cuda.memcpy_dtoh_async(host_outputs[idx], cuda_outputs[idx], stream)127            result[output_name] = host_outputs[idx]128 129        stream.synchronize()130 131        return result132 133 134EXAMPLE_DOC_STRING = """135    Examples:136        ```py137        >>> # !pip install opencv-python transformers accelerate138        >>> from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler139        >>> from diffusers.utils import load_image140        >>> import numpy as np141        >>> import torch142 143        >>> import cv2144        >>> from PIL import Image145 146        >>> # download an image147        >>> image = load_image(148        ...     "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"149        ... )150        >>> np_image = np.array(image)151 152        >>> # get canny image153        >>> np_image = cv2.Canny(np_image, 100, 200)154        >>> np_image = np_image[:, :, None]155        >>> np_image = np.concatenate([np_image, np_image, np_image], axis=2)156        >>> canny_image = Image.fromarray(np_image)157 158        >>> # load control net and stable diffusion v1-5159        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)160        >>> pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(161        ...     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16162        ... )163 164        >>> # speed up diffusion process with faster scheduler and memory optimization165        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)166        >>> pipe.enable_model_cpu_offload()167 168        >>> # generate image169        >>> generator = torch.manual_seed(0)170        >>> image = pipe(171        ...     "futuristic-looking woman",172        ...     num_inference_steps=20,173        ...     generator=generator,174        ...     image=image,175        ...     control_image=canny_image,176        ... ).images[0]177        ```178"""179 180 181def prepare_image(image):182    if isinstance(image, torch.Tensor):183        # Batch single image184        if image.ndim == 3:185            image = image.unsqueeze(0)186 187        image = image.to(dtype=torch.float32)188    else:189        # preprocess image190        if isinstance(image, (PIL.Image.Image, np.ndarray)):191            image = [image]192 193        if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):194            image = [np.array(i.convert("RGB"))[None, :] for i in image]195            image = np.concatenate(image, axis=0)196        elif isinstance(image, list) and isinstance(image[0], np.ndarray):197            image = np.concatenate([i[None, :] for i in image], axis=0)198 199        image = image.transpose(0, 3, 1, 2)200        image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0201 202    return image203 204 205class TensorRTStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):206    vae_encoder: OnnxRuntimeModel207    vae_decoder: OnnxRuntimeModel208    text_encoder: OnnxRuntimeModel209    tokenizer: CLIPTokenizer210    unet: TensorRTModel211    scheduler: KarrasDiffusionSchedulers212 213    def __init__(214        self,215        vae_encoder: OnnxRuntimeModel,216        vae_decoder: OnnxRuntimeModel,217        text_encoder: OnnxRuntimeModel,218        tokenizer: CLIPTokenizer,219        unet: TensorRTModel,220        scheduler: KarrasDiffusionSchedulers,221    ):222        super().__init__()223 224        self.register_modules(225            vae_encoder=vae_encoder,226            vae_decoder=vae_decoder,227            text_encoder=text_encoder,228            tokenizer=tokenizer,229            unet=unet,230            scheduler=scheduler,231        )232        self.vae_scale_factor = 2 ** (4 - 1)233        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)234        self.control_image_processor = VaeImageProcessor(235            vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False236        )237 238    def _encode_prompt(239        self,240        prompt: Union[str, List[str]],241        num_images_per_prompt: Optional[int],242        do_classifier_free_guidance: bool,243        negative_prompt: Optional[str],244        prompt_embeds: Optional[np.ndarray] = None,245        negative_prompt_embeds: Optional[np.ndarray] = None,246    ):247        r"""248        Encodes the prompt into text encoder hidden states.249 250        Args:251            prompt (`str` or `List[str]`):252                prompt to be encoded253            num_images_per_prompt (`int`):254                number of images that should be generated per prompt255            do_classifier_free_guidance (`bool`):256                whether to use classifier free guidance or not257            negative_prompt (`str` or `List[str]`):258                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored259                if `guidance_scale` is less than `1`).260            prompt_embeds (`np.ndarray`, *optional*):261                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not262                provided, text embeddings will be generated from `prompt` input argument.263            negative_prompt_embeds (`np.ndarray`, *optional*):264                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt265                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input266                argument.267        """268        if prompt is not None and isinstance(prompt, str):269            batch_size = 1270        elif prompt is not None and isinstance(prompt, list):271            batch_size = len(prompt)272        else:273            batch_size = prompt_embeds.shape[0]274 275        if prompt_embeds is None:276            # get prompt text embeddings277            text_inputs = self.tokenizer(278                prompt,279                padding="max_length",280                max_length=self.tokenizer.model_max_length,281                truncation=True,282                return_tensors="np",283            )284            text_input_ids = text_inputs.input_ids285            untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids286 287            if not np.array_equal(text_input_ids, untruncated_ids):288                removed_text = self.tokenizer.batch_decode(289                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]290                )291                logger.warning(292                    "The following part of your input was truncated because CLIP can only handle sequences up to"293                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"294                )295 296            prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0]297 298        prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0)299 300        # get unconditional embeddings for classifier free guidance301        if do_classifier_free_guidance and negative_prompt_embeds is None:302            uncond_tokens: List[str]303            if negative_prompt is None:304                uncond_tokens = [""] * batch_size305            elif type(prompt) is not type(negative_prompt):306                raise TypeError(307                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="308                    f" {type(prompt)}."309                )310            elif isinstance(negative_prompt, str):311                uncond_tokens = [negative_prompt] * batch_size312            elif batch_size != len(negative_prompt):313                raise ValueError(314                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"315                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"316                    " the batch size of `prompt`."317                )318            else:319                uncond_tokens = negative_prompt320 321            max_length = prompt_embeds.shape[1]322            uncond_input = self.tokenizer(323                uncond_tokens,324                padding="max_length",325                max_length=max_length,326                truncation=True,327                return_tensors="np",328            )329            negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0]330 331        if do_classifier_free_guidance:332            negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0)333 334            # For classifier free guidance, we need to do two forward passes.335            # Here we concatenate the unconditional and text embeddings into a single batch336            # to avoid doing two forward passes337            prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds])338 339        return prompt_embeds340 341    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents342    def decode_latents(self, latents):343        warnings.warn(344            "The decode_latents method is deprecated and will be removed in a future version. Please"345            " use VaeImageProcessor instead",346            FutureWarning,347        )348        latents = 1 / self.vae.config.scaling_factor * latents349        image = self.vae.decode(latents, return_dict=False)[0]350        image = (image / 2 + 0.5).clamp(0, 1)351        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16352        image = image.cpu().permute(0, 2, 3, 1).float().numpy()353        return image354 355    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs356    def prepare_extra_step_kwargs(self, generator, eta):357        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature358        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.359        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502360        # and should be between [0, 1]361 362        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())363        extra_step_kwargs = {}364        if accepts_eta:365            extra_step_kwargs["eta"] = eta366 367        # check if the scheduler accepts generator368        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())369        if accepts_generator:370            extra_step_kwargs["generator"] = generator371        return extra_step_kwargs372 373    def check_inputs(374        self,375        num_controlnet,376        prompt,377        image,378        callback_steps,379        negative_prompt=None,380        prompt_embeds=None,381        negative_prompt_embeds=None,382        controlnet_conditioning_scale=1.0,383        control_guidance_start=0.0,384        control_guidance_end=1.0,385    ):386        if (callback_steps is None) or (387            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)388        ):389            raise ValueError(390                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"391                f" {type(callback_steps)}."392            )393 394        if prompt is not None and prompt_embeds is not None:395            raise ValueError(396                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"397                " only forward one of the two."398            )399        elif prompt is None and prompt_embeds is None:400            raise ValueError(401                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."402            )403        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):404            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")405 406        if negative_prompt is not None and negative_prompt_embeds is not None:407            raise ValueError(408                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"409                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."410            )411 412        if prompt_embeds is not None and negative_prompt_embeds is not None:413            if prompt_embeds.shape != negative_prompt_embeds.shape:414                raise ValueError(415                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"416                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"417                    f" {negative_prompt_embeds.shape}."418                )419 420        # Check `image`421        if num_controlnet == 1:422            self.check_image(image, prompt, prompt_embeds)423        elif num_controlnet > 1:424            if not isinstance(image, list):425                raise TypeError("For multiple controlnets: `image` must be type `list`")426 427            # When `image` is a nested list:428            # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])429            elif any(isinstance(i, list) for i in image):430                raise ValueError("A single batch of multiple conditionings are supported at the moment.")431            elif len(image) != num_controlnet:432                raise ValueError(433                    f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {num_controlnet} ControlNets."434                )435 436            for image_ in image:437                self.check_image(image_, prompt, prompt_embeds)438        else:439            assert False440 441        # Check `controlnet_conditioning_scale`442        if num_controlnet == 1:443            if not isinstance(controlnet_conditioning_scale, float):444                raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")445        elif num_controlnet > 1:446            if isinstance(controlnet_conditioning_scale, list):447                if any(isinstance(i, list) for i in controlnet_conditioning_scale):448                    raise ValueError("A single batch of multiple conditionings are supported at the moment.")449            elif (450                isinstance(controlnet_conditioning_scale, list)451                and len(controlnet_conditioning_scale) != num_controlnet452            ):453                raise ValueError(454                    "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"455                    " the same length as the number of controlnets"456                )457        else:458            assert False459 460        if len(control_guidance_start) != len(control_guidance_end):461            raise ValueError(462                f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list."463            )464 465        if num_controlnet > 1:466            if len(control_guidance_start) != num_controlnet:467                raise ValueError(468                    f"`control_guidance_start`: {control_guidance_start} has {len(control_guidance_start)} elements but there are {num_controlnet} controlnets available. Make sure to provide {num_controlnet}."469                )470 471        for start, end in zip(control_guidance_start, control_guidance_end):472            if start >= end:473                raise ValueError(474                    f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."475                )476            if start < 0.0:477                raise ValueError(f"control guidance start: {start} can't be smaller than 0.")478            if end > 1.0:479                raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")480 481    # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.check_image482    def check_image(self, image, prompt, prompt_embeds):483        image_is_pil = isinstance(image, PIL.Image.Image)484        image_is_tensor = isinstance(image, torch.Tensor)485        image_is_np = isinstance(image, np.ndarray)486        image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)487        image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)488        image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)489 490        if (491            not image_is_pil492            and not image_is_tensor493            and not image_is_np494            and not image_is_pil_list495            and not image_is_tensor_list496            and not image_is_np_list497        ):498            raise TypeError(499                f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}"500            )501 502        if image_is_pil:503            image_batch_size = 1504        else:505            image_batch_size = len(image)506 507        if prompt is not None and isinstance(prompt, str):508            prompt_batch_size = 1509        elif prompt is not None and isinstance(prompt, list):510            prompt_batch_size = len(prompt)511        elif prompt_embeds is not None:512            prompt_batch_size = prompt_embeds.shape[0]513 514        if image_batch_size != 1 and image_batch_size != prompt_batch_size:515            raise ValueError(516                f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"517            )518 519    # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image520    def prepare_control_image(521        self,522        image,523        width,524        height,525        batch_size,526        num_images_per_prompt,527        device,528        dtype,529        do_classifier_free_guidance=False,530        guess_mode=False,531    ):532        image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)533        image_batch_size = image.shape[0]534 535        if image_batch_size == 1:536            repeat_by = batch_size537        else:538            # image batch size is the same as prompt batch size539            repeat_by = num_images_per_prompt540 541        image = image.repeat_interleave(repeat_by, dim=0)542 543        image = image.to(device=device, dtype=dtype)544 545        if do_classifier_free_guidance and not guess_mode:546            image = torch.cat([image] * 2)547 548        return image549 550    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps551    def get_timesteps(self, num_inference_steps, strength, device):552        # get the original timestep using init_timestep553        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)554 555        t_start = max(num_inference_steps - init_timestep, 0)556        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]557 558        return timesteps, num_inference_steps - t_start559 560    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):561        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):562            raise ValueError(563                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"564            )565 566        image = image.to(device=device, dtype=dtype)567 568        batch_size = batch_size * num_images_per_prompt569 570        if image.shape[1] == 4:571            init_latents = image572 573        else:574            _image = image.cpu().detach().numpy()575            init_latents = self.vae_encoder(sample=_image)[0]576            init_latents = torch.from_numpy(init_latents).to(device=device, dtype=dtype)577            init_latents = 0.18215 * init_latents578 579        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:580            # expand init_latents for batch_size581            deprecation_message = (582                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"583                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"584                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"585                " your script to pass as many initial images as text prompts to suppress this warning."586            )587            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)588            additional_image_per_prompt = batch_size // init_latents.shape[0]589            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)590        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:591            raise ValueError(592                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."593            )594        else:595            init_latents = torch.cat([init_latents], dim=0)596 597        shape = init_latents.shape598        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)599 600        # get latents601        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)602        latents = init_latents603 604        return latents605 606    @torch.no_grad()607    @replace_example_docstring(EXAMPLE_DOC_STRING)608    def __call__(609        self,610        num_controlnet: int,611        fp16: bool = True,612        prompt: Union[str, List[str]] = None,613        image: Union[614            torch.Tensor,615            PIL.Image.Image,616            np.ndarray,617            List[torch.Tensor],618            List[PIL.Image.Image],619            List[np.ndarray],620        ] = None,621        control_image: Union[622            torch.Tensor,623            PIL.Image.Image,624            np.ndarray,625            List[torch.Tensor],626            List[PIL.Image.Image],627            List[np.ndarray],628        ] = None,629        height: Optional[int] = None,630        width: Optional[int] = None,631        strength: float = 0.8,632        num_inference_steps: int = 50,633        guidance_scale: float = 7.5,634        negative_prompt: Optional[Union[str, List[str]]] = None,635        num_images_per_prompt: Optional[int] = 1,636        eta: float = 0.0,637        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,638        latents: Optional[torch.Tensor] = None,639        prompt_embeds: Optional[torch.Tensor] = None,640        negative_prompt_embeds: Optional[torch.Tensor] = None,641        output_type: Optional[str] = "pil",642        return_dict: bool = True,643        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,644        callback_steps: int = 1,645        cross_attention_kwargs: Optional[Dict[str, Any]] = None,646        controlnet_conditioning_scale: Union[float, List[float]] = 0.8,647        guess_mode: bool = False,648        control_guidance_start: Union[float, List[float]] = 0.0,649        control_guidance_end: Union[float, List[float]] = 1.0,650    ):651        r"""652        Function invoked when calling the pipeline for generation.653 654        Args:655            prompt (`str` or `List[str]`, *optional*):656                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.657                instead.658            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:659                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):660                The initial image will be used as the starting point for the image generation process. Can also accept661                image latents as `image`, if passing latents directly, it will not be encoded again.662            control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:663                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):664                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If665                the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can666                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If667                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are668                specified in init, images must be passed as a list such that each element of the list can be correctly669                batched for input to a single controlnet.670            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):671                The height in pixels of the generated image.672            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):673                The width in pixels of the generated image.674            num_inference_steps (`int`, *optional*, defaults to 50):675                The number of denoising steps. More denoising steps usually lead to a higher quality image at the676                expense of slower inference.677            guidance_scale (`float`, *optional*, defaults to 7.5):678                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).679                `guidance_scale` is defined as `w` of equation 2. of [Imagen680                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >681                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,682                usually at the expense of lower image quality.683            negative_prompt (`str` or `List[str]`, *optional*):684                The prompt or prompts not to guide the image generation. If not defined, one has to pass685                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is686                less than `1`).687            num_images_per_prompt (`int`, *optional*, defaults to 1):688                The number of images to generate per prompt.689            eta (`float`, *optional*, defaults to 0.0):690                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to691                [`schedulers.DDIMScheduler`], will be ignored for others.692            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):693                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)694                to make generation deterministic.695            latents (`torch.Tensor`, *optional*):696                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image697                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents698                tensor will ge generated by sampling using the supplied random `generator`.699            prompt_embeds (`torch.Tensor`, *optional*):700                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not701                provided, text embeddings will be generated from `prompt` input argument.702            negative_prompt_embeds (`torch.Tensor`, *optional*):703                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt704                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input705                argument.706            output_type (`str`, *optional*, defaults to `"pil"`):707                The output format of the generate image. Choose between708                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.709            return_dict (`bool`, *optional*, defaults to `True`):710                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a711                plain tuple.712            callback (`Callable`, *optional*):713                A function that will be called every `callback_steps` steps during inference. The function will be714                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.715            callback_steps (`int`, *optional*, defaults to 1):716                The frequency at which the `callback` function will be called. If not specified, the callback will be717                called at every step.718            cross_attention_kwargs (`dict`, *optional*):719                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under720                `self.processor` in721                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).722            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):723                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added724                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the725                corresponding scale as a list. Note that by default, we use a smaller conditioning scale for inpainting726                than for [`~StableDiffusionControlNetPipeline.__call__`].727            guess_mode (`bool`, *optional*, defaults to `False`):728                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if729                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.730            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):731                The percentage of total steps at which the controlnet starts applying.732            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):733                The percentage of total steps at which the controlnet stops applying.734 735        Examples:736 737        Returns:738            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:739            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.740            When returning a tuple, the first element is a list with the generated images, and the second element is a741            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"742            (nsfw) content, according to the `safety_checker`.743        """744        if fp16:745            torch_dtype = torch.float16746            np_dtype = np.float16747        else:748            torch_dtype = torch.float32749            np_dtype = np.float32750 751        # align format for control guidance752        if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):753            control_guidance_start = len(control_guidance_end) * [control_guidance_start]754        elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):755            control_guidance_end = len(control_guidance_start) * [control_guidance_end]756        elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):757            mult = num_controlnet758            control_guidance_start, control_guidance_end = (759                mult * [control_guidance_start],760                mult * [control_guidance_end],761            )762 763        # 1. Check inputs. Raise error if not correct764        self.check_inputs(765            num_controlnet,766            prompt,767            control_image,768            callback_steps,769            negative_prompt,770            prompt_embeds,771            negative_prompt_embeds,772            controlnet_conditioning_scale,773            control_guidance_start,774            control_guidance_end,775        )776 777        # 2. Define call parameters778        if prompt is not None and isinstance(prompt, str):779            batch_size = 1780        elif prompt is not None and isinstance(prompt, list):781            batch_size = len(prompt)782        else:783            batch_size = prompt_embeds.shape[0]784 785        device = self._execution_device786        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)787        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`788        # corresponds to doing no classifier free guidance.789        do_classifier_free_guidance = guidance_scale > 1.0790 791        if num_controlnet > 1 and isinstance(controlnet_conditioning_scale, float):792            controlnet_conditioning_scale = [controlnet_conditioning_scale] * num_controlnet793 794        # 3. Encode input prompt795        prompt_embeds = self._encode_prompt(796            prompt,797            num_images_per_prompt,798            do_classifier_free_guidance,799            negative_prompt,800            prompt_embeds=prompt_embeds,801            negative_prompt_embeds=negative_prompt_embeds,802        )803        # 4. Prepare image804        image = self.image_processor.preprocess(image).to(dtype=torch.float32)805 806        # 5. Prepare controlnet_conditioning_image807        if num_controlnet == 1:808            control_image = self.prepare_control_image(809                image=control_image,810                width=width,811                height=height,812                batch_size=batch_size * num_images_per_prompt,813                num_images_per_prompt=num_images_per_prompt,814                device=device,815                dtype=torch_dtype,816                do_classifier_free_guidance=do_classifier_free_guidance,817                guess_mode=guess_mode,818            )819        elif num_controlnet > 1:820            control_images = []821 822            for control_image_ in control_image:823                control_image_ = self.prepare_control_image(824                    image=control_image_,825                    width=width,826                    height=height,827                    batch_size=batch_size * num_images_per_prompt,828                    num_images_per_prompt=num_images_per_prompt,829                    device=device,830                    dtype=torch_dtype,831                    do_classifier_free_guidance=do_classifier_free_guidance,832                    guess_mode=guess_mode,833                )834 835                control_images.append(control_image_)836 837            control_image = control_images838        else:839            assert False840 841        # 5. Prepare timesteps842        self.scheduler.set_timesteps(num_inference_steps, device=device)843        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)844        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)845 846        # 6. Prepare latent variables847        latents = self.prepare_latents(848            image,849            latent_timestep,850            batch_size,851            num_images_per_prompt,852            torch_dtype,853            device,854            generator,855        )856 857        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline858        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)859 860        # 7.1 Create tensor stating which controlnets to keep861        controlnet_keep = []862        for i in range(len(timesteps)):863            keeps = [864                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)865                for s, e in zip(control_guidance_start, control_guidance_end)866            ]867            controlnet_keep.append(keeps[0] if num_controlnet == 1 else keeps)868 869        # 8. Denoising loop870        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order871        with self.progress_bar(total=num_inference_steps) as progress_bar:872            for i, t in enumerate(timesteps):873                # expand the latents if we are doing classifier free guidance874                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents875                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)876 877                if isinstance(controlnet_keep[i], list):878                    cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]879                else:880                    controlnet_cond_scale = controlnet_conditioning_scale881                    if isinstance(controlnet_cond_scale, list):882                        controlnet_cond_scale = controlnet_cond_scale[0]883                    cond_scale = controlnet_cond_scale * controlnet_keep[i]884 885                # predict the noise residual886                _latent_model_input = latent_model_input.cpu().detach().numpy()887                _prompt_embeds = np.array(prompt_embeds, dtype=np_dtype)888                _t = np.array([t.cpu().detach().numpy()], dtype=np_dtype)889 890                if num_controlnet == 1:891                    control_images = np.array([control_image], dtype=np_dtype)892                else:893                    control_images = []894                    for _control_img in control_image:895                        _control_img = _control_img.cpu().detach().numpy()896                        control_images.append(_control_img)897                    control_images = np.array(control_images, dtype=np_dtype)898 899                control_scales = np.array(cond_scale, dtype=np_dtype)900                control_scales = np.resize(control_scales, (num_controlnet, 1))901 902                noise_pred = self.unet(903                    sample=_latent_model_input,904                    timestep=_t,905                    encoder_hidden_states=_prompt_embeds,906                    controlnet_conds=control_images,907                    conditioning_scales=control_scales,908                )["noise_pred"]909                noise_pred = torch.from_numpy(noise_pred).to(device)910 911                # perform guidance912                if do_classifier_free_guidance:913                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)914                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)915 916                # compute the previous noisy sample x_t -> x_t-1917                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]918 919                # call the callback, if provided920                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):921                    progress_bar.update()922                    if callback is not None and i % callback_steps == 0:923                        step_idx = i // getattr(self.scheduler, "order", 1)924                        callback(step_idx, t, latents)925 926        if not output_type == "latent":927            _latents = latents.cpu().detach().numpy() / 0.18215928            _latents = np.array(_latents, dtype=np_dtype)929            image = self.vae_decoder(latent_sample=_latents)[0]930            image = torch.from_numpy(image).to(device, dtype=torch.float32)931            has_nsfw_concept = None932        else:933            image = latents934            has_nsfw_concept = None935 936        if has_nsfw_concept is None:937            do_denormalize = [True] * image.shape[0]938        else:939            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]940 941        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)942 943        if not return_dict:944            return (image, has_nsfw_concept)945 946        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)947 948 949if __name__ == "__main__":950    parser = argparse.ArgumentParser()951 952    parser.add_argument(953        "--sd_model",954        type=str,955        required=True,956        help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",957    )958 959    parser.add_argument(960        "--onnx_model_dir",961        type=str,962        required=True,963        help="Path to the ONNX directory",964    )965 966    parser.add_argument(967        "--unet_engine_path",968        type=str,969        required=True,970        help="Path to the unet + controlnet tensorrt model",971    )972 973    parser.add_argument("--qr_img_path", type=str, required=True, help="Path to the qr code image")974 975    args = parser.parse_args()976 977    qr_image = Image.open(args.qr_img_path)978    qr_image = qr_image.resize((512, 512))979 980    # init stable diffusion pipeline981    pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(args.sd_model)982    pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)983 984    provider = ["CUDAExecutionProvider", "CPUExecutionProvider"]985    onnx_pipeline = TensorRTStableDiffusionControlNetImg2ImgPipeline(986        vae_encoder=OnnxRuntimeModel.from_pretrained(987            os.path.join(args.onnx_model_dir, "vae_encoder"), provider=provider988        ),989        vae_decoder=OnnxRuntimeModel.from_pretrained(990            os.path.join(args.onnx_model_dir, "vae_decoder"), provider=provider991        ),992        text_encoder=OnnxRuntimeModel.from_pretrained(993            os.path.join(args.onnx_model_dir, "text_encoder"), provider=provider994        ),995        tokenizer=pipeline.tokenizer,996        unet=TensorRTModel(args.unet_engine_path),997        scheduler=pipeline.scheduler,998    )999    onnx_pipeline = onnx_pipeline.to("cuda")1000 1001    prompt = "a cute cat fly to the moon"1002    negative_prompt = "paintings, sketches, worst quality, low quality, normal quality, lowres, normal quality, monochrome, grayscale, skin spots, acnes, skin blemishes, age spot, glans, nsfw, nipples, necklace, worst quality, low quality, watermark, username, signature, multiple breasts, lowres, bad anatomy, bad hands, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad feet, single color, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, ugly, blurry, bad anatomy, bad proportions, extra limbs, disfigured, bad anatomy, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, mutated hands, fused fingers, too many fingers, long neck, bad body perspect"1003 1004    for i in range(10):1005        start_time = time.time()1006        image = onnx_pipeline(1007            num_controlnet=2,1008            prompt=prompt,1009            negative_prompt=negative_prompt,1010            image=qr_image,1011            control_image=[qr_image, qr_image],1012            width=512,1013            height=512,1014            strength=0.75,1015            num_inference_steps=20,1016            num_images_per_prompt=1,1017            controlnet_conditioning_scale=[0.8, 0.8],1018            control_guidance_start=[0.3, 0.3],1019            control_guidance_end=[0.9, 0.9],1020        ).images[0]1021        print(time.time() - start_time)1022        image.save("output_qr_code.png")1023