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    randn_tensor,27    replace_example_docstring,28)29 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.FloatTensor,615            PIL.Image.Image,616            np.ndarray,617            List[torch.FloatTensor],618            List[PIL.Image.Image],619            List[np.ndarray],620        ] = None,621        control_image: Union[622            torch.FloatTensor,623            PIL.Image.Image,624            np.ndarray,625            List[torch.FloatTensor],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.FloatTensor] = None,639        prompt_embeds: Optional[torch.FloatTensor] = None,640        negative_prompt_embeds: Optional[torch.FloatTensor] = None,641        output_type: Optional[str] = "pil",642        return_dict: bool = True,643        callback: Optional[Callable[[int, int, torch.FloatTensor], 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.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:659                    `List[List[torch.FloatTensor]]`, `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 accpet661                image latents as `image`, if passing latents directly, it will not be encoded again.662            control_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:663                    `List[List[torch.FloatTensor]]`, `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.FloatTensor`, 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.FloatTensor`, *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.FloatTensor`, *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.FloatTensor`, *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.FloatTensor)`.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 = mult * [control_guidance_start], mult * [759                control_guidance_end760            ]761 762        # 1. Check inputs. Raise error if not correct763        self.check_inputs(764            num_controlnet,765            prompt,766            control_image,767            callback_steps,768            negative_prompt,769            prompt_embeds,770            negative_prompt_embeds,771            controlnet_conditioning_scale,772            control_guidance_start,773            control_guidance_end,774        )775 776        # 2. Define call parameters777        if prompt is not None and isinstance(prompt, str):778            batch_size = 1779        elif prompt is not None and isinstance(prompt, list):780            batch_size = len(prompt)781        else:782            batch_size = prompt_embeds.shape[0]783 784        device = self._execution_device785        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)786        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`787        # corresponds to doing no classifier free guidance.788        do_classifier_free_guidance = guidance_scale > 1.0789 790        if num_controlnet > 1 and isinstance(controlnet_conditioning_scale, float):791            controlnet_conditioning_scale = [controlnet_conditioning_scale] * num_controlnet792 793        # 3. Encode input prompt794        prompt_embeds = self._encode_prompt(795            prompt,796            num_images_per_prompt,797            do_classifier_free_guidance,798            negative_prompt,799            prompt_embeds=prompt_embeds,800            negative_prompt_embeds=negative_prompt_embeds,801        )802        # 4. Prepare image803        image = self.image_processor.preprocess(image).to(dtype=torch.float32)804 805        # 5. Prepare controlnet_conditioning_image806        if num_controlnet == 1:807            control_image = self.prepare_control_image(808                image=control_image,809                width=width,810                height=height,811                batch_size=batch_size * num_images_per_prompt,812                num_images_per_prompt=num_images_per_prompt,813                device=device,814                dtype=torch_dtype,815                do_classifier_free_guidance=do_classifier_free_guidance,816                guess_mode=guess_mode,817            )818        elif num_controlnet > 1:819            control_images = []820 821            for control_image_ in control_image:822                control_image_ = self.prepare_control_image(823                    image=control_image_,824                    width=width,825                    height=height,826                    batch_size=batch_size * num_images_per_prompt,827                    num_images_per_prompt=num_images_per_prompt,828                    device=device,829                    dtype=torch_dtype,830                    do_classifier_free_guidance=do_classifier_free_guidance,831                    guess_mode=guess_mode,832                )833 834                control_images.append(control_image_)835 836            control_image = control_images837        else:838            assert False839 840        # 5. Prepare timesteps841        self.scheduler.set_timesteps(num_inference_steps, device=device)842        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)843        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)844 845        # 6. Prepare latent variables846        latents = self.prepare_latents(847            image,848            latent_timestep,849            batch_size,850            num_images_per_prompt,851            torch_dtype,852            device,853            generator,854        )855 856        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline857        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)858 859        # 7.1 Create tensor stating which controlnets to keep860        controlnet_keep = []861        for i in range(len(timesteps)):862            keeps = [863                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)864                for s, e in zip(control_guidance_start, control_guidance_end)865            ]866            controlnet_keep.append(keeps[0] if num_controlnet == 1 else keeps)867 868        # 8. Denoising loop869        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order870        with self.progress_bar(total=num_inference_steps) as progress_bar:871            for i, t in enumerate(timesteps):872                # expand the latents if we are doing classifier free guidance873                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents874                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)875 876                if isinstance(controlnet_keep[i], list):877                    cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]878                else:879                    controlnet_cond_scale = controlnet_conditioning_scale880                    if isinstance(controlnet_cond_scale, list):881                        controlnet_cond_scale = controlnet_cond_scale[0]882                    cond_scale = controlnet_cond_scale * controlnet_keep[i]883 884                # predict the noise residual885                _latent_model_input = latent_model_input.cpu().detach().numpy()886                _prompt_embeds = np.array(prompt_embeds, dtype=np_dtype)887                _t = np.array([t.cpu().detach().numpy()], dtype=np_dtype)888 889                if num_controlnet == 1:890                    control_images = np.array([control_image], dtype=np_dtype)891                else:892                    control_images = []893                    for _control_img in control_image:894                        _control_img = _control_img.cpu().detach().numpy()895                        control_images.append(_control_img)896                    control_images = np.array(control_images, dtype=np_dtype)897 898                control_scales = np.array(cond_scale, dtype=np_dtype)899                control_scales = np.resize(control_scales, (num_controlnet, 1))900 901                noise_pred = self.unet(902                    sample=_latent_model_input,903                    timestep=_t,904                    encoder_hidden_states=_prompt_embeds,905                    controlnet_conds=control_images,906                    conditioning_scales=control_scales,907                )["noise_pred"]908                noise_pred = torch.from_numpy(noise_pred).to(device)909 910                # perform guidance911                if do_classifier_free_guidance:912                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)913                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)914 915                # compute the previous noisy sample x_t -> x_t-1916                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]917 918                # call the callback, if provided919                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):920                    progress_bar.update()921                    if callback is not None and i % callback_steps == 0:922                        callback(i, t, latents)923 924        if not output_type == "latent":925            _latents = latents.cpu().detach().numpy() / 0.18215926            _latents = np.array(_latents, dtype=np_dtype)927            image = self.vae_decoder(latent_sample=_latents)[0]928            image = torch.from_numpy(image).to(device, dtype=torch.float32)929            has_nsfw_concept = None930        else:931            image = latents932            has_nsfw_concept = None933 934        if has_nsfw_concept is None:935            do_denormalize = [True] * image.shape[0]936        else:937            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]938 939        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)940 941        if not return_dict:942            return (image, has_nsfw_concept)943 944        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)945 946 947if __name__ == "__main__":948    parser = argparse.ArgumentParser()949 950    parser.add_argument(951        "--sd_model",952        type=str,953        required=True,954        help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",955    )956 957    parser.add_argument(958        "--onnx_model_dir",959        type=str,960        required=True,961        help="Path to the ONNX directory",962    )963 964    parser.add_argument(965        "--unet_engine_path",966        type=str,967        required=True,968        help="Path to the unet + controlnet tensorrt model",969    )970 971    parser.add_argument("--qr_img_path", type=str, required=True, help="Path to the qr code image")972 973    args = parser.parse_args()974 975    qr_image = Image.open(args.qr_img_path)976    qr_image = qr_image.resize((512, 512))977 978    # init stable diffusion pipeline979    pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(args.sd_model)980    pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)981 982    provider = ["CUDAExecutionProvider", "CPUExecutionProvider"]983    onnx_pipeline = TensorRTStableDiffusionControlNetImg2ImgPipeline(984        vae_encoder=OnnxRuntimeModel.from_pretrained(985            os.path.join(args.onnx_model_dir, "vae_encoder"), provider=provider986        ),987        vae_decoder=OnnxRuntimeModel.from_pretrained(988            os.path.join(args.onnx_model_dir, "vae_decoder"), provider=provider989        ),990        text_encoder=OnnxRuntimeModel.from_pretrained(991            os.path.join(args.onnx_model_dir, "text_encoder"), provider=provider992        ),993        tokenizer=pipeline.tokenizer,994        unet=TensorRTModel(args.unet_engine_path),995        scheduler=pipeline.scheduler,996    )997    onnx_pipeline = onnx_pipeline.to("cuda")998 999    prompt = "a cute cat fly to the moon"1000    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, tranny, trans, trannsexual, hermaphrodite, 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"1001 1002    for i in range(10):1003        start_time = time.time()1004        image = onnx_pipeline(1005            num_controlnet=2,1006            prompt=prompt,1007            negative_prompt=negative_prompt,1008            image=qr_image,1009            control_image=[qr_image, qr_image],1010            width=512,1011            height=512,1012            strength=0.75,1013            num_inference_steps=20,1014            num_images_per_prompt=1,1015            controlnet_conditioning_scale=[0.8, 0.8],1016            control_guidance_start=[0.3, 0.3],1017            control_guidance_end=[0.9, 0.9],1018        ).images[0]1019        print(time.time() - start_time)1020        image.save("output_qr_code.png")1021