diffusers/community-pipelines-mirror

1import argparse2import inspect3import os4import time5import warnings6from typing import Any, Callable, Dict, List, Optional, Union7 8import numpy as np9import PIL.Image10import torch11from PIL import Image12from transformers import CLIPTokenizer13 14from diffusers import OnnxRuntimeModel, StableDiffusionImg2ImgPipeline, UniPCMultistepScheduler15from diffusers.image_processor import VaeImageProcessor16from diffusers.pipelines.pipeline_utils import DiffusionPipeline17from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput18from diffusers.schedulers import KarrasDiffusionSchedulers19from diffusers.utils import (20    deprecate,21    logging,22    replace_example_docstring,23)24from diffusers.utils.torch_utils import randn_tensor25 26 27logger = logging.get_logger(__name__)  # pylint: disable=invalid-name28 29 30EXAMPLE_DOC_STRING = """31    Examples:32        ```py33        >>> # !pip install opencv-python transformers accelerate34        >>> from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler35        >>> from diffusers.utils import load_image36        >>> import numpy as np37        >>> import torch38 39        >>> import cv240        >>> from PIL import Image41 42        >>> # download an image43        >>> image = load_image(44        ...     "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"45        ... )46        >>> np_image = np.array(image)47 48        >>> # get canny image49        >>> np_image = cv2.Canny(np_image, 100, 200)50        >>> np_image = np_image[:, :, None]51        >>> np_image = np.concatenate([np_image, np_image, np_image], axis=2)52        >>> canny_image = Image.fromarray(np_image)53 54        >>> # load control net and stable diffusion v1-555        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)56        >>> pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(57        ...     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float1658        ... )59 60        >>> # speed up diffusion process with faster scheduler and memory optimization61        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)62        >>> pipe.enable_model_cpu_offload()63 64        >>> # generate image65        >>> generator = torch.manual_seed(0)66        >>> image = pipe(67        ...     "futuristic-looking woman",68        ...     num_inference_steps=20,69        ...     generator=generator,70        ...     image=image,71        ...     control_image=canny_image,72        ... ).images[0]73        ```74"""75 76 77def prepare_image(image):78    if isinstance(image, torch.Tensor):79        # Batch single image80        if image.ndim == 3:81            image = image.unsqueeze(0)82 83        image = image.to(dtype=torch.float32)84    else:85        # preprocess image86        if isinstance(image, (PIL.Image.Image, np.ndarray)):87            image = [image]88 89        if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):90            image = [np.array(i.convert("RGB"))[None, :] for i in image]91            image = np.concatenate(image, axis=0)92        elif isinstance(image, list) and isinstance(image[0], np.ndarray):93            image = np.concatenate([i[None, :] for i in image], axis=0)94 95        image = image.transpose(0, 3, 1, 2)96        image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.097 98    return image99 100 101class OnnxStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):102    vae_encoder: OnnxRuntimeModel103    vae_decoder: OnnxRuntimeModel104    text_encoder: OnnxRuntimeModel105    tokenizer: CLIPTokenizer106    unet: OnnxRuntimeModel107    scheduler: KarrasDiffusionSchedulers108 109    def __init__(110        self,111        vae_encoder: OnnxRuntimeModel,112        vae_decoder: OnnxRuntimeModel,113        text_encoder: OnnxRuntimeModel,114        tokenizer: CLIPTokenizer,115        unet: OnnxRuntimeModel,116        scheduler: KarrasDiffusionSchedulers,117    ):118        super().__init__()119 120        self.register_modules(121            vae_encoder=vae_encoder,122            vae_decoder=vae_decoder,123            text_encoder=text_encoder,124            tokenizer=tokenizer,125            unet=unet,126            scheduler=scheduler,127        )128        self.vae_scale_factor = 2 ** (4 - 1)129        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)130        self.control_image_processor = VaeImageProcessor(131            vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False132        )133 134    def _encode_prompt(135        self,136        prompt: Union[str, List[str]],137        num_images_per_prompt: Optional[int],138        do_classifier_free_guidance: bool,139        negative_prompt: Optional[str],140        prompt_embeds: Optional[np.ndarray] = None,141        negative_prompt_embeds: Optional[np.ndarray] = None,142    ):143        r"""144        Encodes the prompt into text encoder hidden states.145 146        Args:147            prompt (`str` or `List[str]`):148                prompt to be encoded149            num_images_per_prompt (`int`):150                number of images that should be generated per prompt151            do_classifier_free_guidance (`bool`):152                whether to use classifier free guidance or not153            negative_prompt (`str` or `List[str]`):154                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored155                if `guidance_scale` is less than `1`).156            prompt_embeds (`np.ndarray`, *optional*):157                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not158                provided, text embeddings will be generated from `prompt` input argument.159            negative_prompt_embeds (`np.ndarray`, *optional*):160                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt161                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input162                argument.163        """164        if prompt is not None and isinstance(prompt, str):165            batch_size = 1166        elif prompt is not None and isinstance(prompt, list):167            batch_size = len(prompt)168        else:169            batch_size = prompt_embeds.shape[0]170 171        if prompt_embeds is None:172            # get prompt text embeddings173            text_inputs = self.tokenizer(174                prompt,175                padding="max_length",176                max_length=self.tokenizer.model_max_length,177                truncation=True,178                return_tensors="np",179            )180            text_input_ids = text_inputs.input_ids181            untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids182 183            if not np.array_equal(text_input_ids, untruncated_ids):184                removed_text = self.tokenizer.batch_decode(185                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]186                )187                logger.warning(188                    "The following part of your input was truncated because CLIP can only handle sequences up to"189                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"190                )191 192            prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0]193 194        prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0)195 196        # get unconditional embeddings for classifier free guidance197        if do_classifier_free_guidance and negative_prompt_embeds is None:198            uncond_tokens: List[str]199            if negative_prompt is None:200                uncond_tokens = [""] * batch_size201            elif type(prompt) is not type(negative_prompt):202                raise TypeError(203                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="204                    f" {type(prompt)}."205                )206            elif isinstance(negative_prompt, str):207                uncond_tokens = [negative_prompt] * batch_size208            elif batch_size != len(negative_prompt):209                raise ValueError(210                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"211                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"212                    " the batch size of `prompt`."213                )214            else:215                uncond_tokens = negative_prompt216 217            max_length = prompt_embeds.shape[1]218            uncond_input = self.tokenizer(219                uncond_tokens,220                padding="max_length",221                max_length=max_length,222                truncation=True,223                return_tensors="np",224            )225            negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0]226 227        if do_classifier_free_guidance:228            negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0)229 230            # For classifier free guidance, we need to do two forward passes.231            # Here we concatenate the unconditional and text embeddings into a single batch232            # to avoid doing two forward passes233            prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds])234 235        return prompt_embeds236 237    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents238    def decode_latents(self, latents):239        warnings.warn(240            "The decode_latents method is deprecated and will be removed in a future version. Please"241            " use VaeImageProcessor instead",242            FutureWarning,243        )244        latents = 1 / self.vae.config.scaling_factor * latents245        image = self.vae.decode(latents, return_dict=False)[0]246        image = (image / 2 + 0.5).clamp(0, 1)247        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16248        image = image.cpu().permute(0, 2, 3, 1).float().numpy()249        return image250 251    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs252    def prepare_extra_step_kwargs(self, generator, eta):253        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature254        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.255        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502256        # and should be between [0, 1]257 258        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())259        extra_step_kwargs = {}260        if accepts_eta:261            extra_step_kwargs["eta"] = eta262 263        # check if the scheduler accepts generator264        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())265        if accepts_generator:266            extra_step_kwargs["generator"] = generator267        return extra_step_kwargs268 269    def check_inputs(270        self,271        num_controlnet,272        prompt,273        image,274        callback_steps,275        negative_prompt=None,276        prompt_embeds=None,277        negative_prompt_embeds=None,278        controlnet_conditioning_scale=1.0,279        control_guidance_start=0.0,280        control_guidance_end=1.0,281    ):282        if (callback_steps is None) or (283            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)284        ):285            raise ValueError(286                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"287                f" {type(callback_steps)}."288            )289 290        if prompt is not None and prompt_embeds is not None:291            raise ValueError(292                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"293                " only forward one of the two."294            )295        elif prompt is None and prompt_embeds is None:296            raise ValueError(297                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."298            )299        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):300            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")301 302        if negative_prompt is not None and negative_prompt_embeds is not None:303            raise ValueError(304                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"305                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."306            )307 308        if prompt_embeds is not None and negative_prompt_embeds is not None:309            if prompt_embeds.shape != negative_prompt_embeds.shape:310                raise ValueError(311                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"312                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"313                    f" {negative_prompt_embeds.shape}."314                )315 316        # Check `image`317        if num_controlnet == 1:318            self.check_image(image, prompt, prompt_embeds)319        elif num_controlnet > 1:320            if not isinstance(image, list):321                raise TypeError("For multiple controlnets: `image` must be type `list`")322 323            # When `image` is a nested list:324            # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])325            elif any(isinstance(i, list) for i in image):326                raise ValueError("A single batch of multiple conditionings are supported at the moment.")327            elif len(image) != num_controlnet:328                raise ValueError(329                    f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {num_controlnet} ControlNets."330                )331 332            for image_ in image:333                self.check_image(image_, prompt, prompt_embeds)334        else:335            assert False336 337        # Check `controlnet_conditioning_scale`338        if num_controlnet == 1:339            if not isinstance(controlnet_conditioning_scale, float):340                raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")341        elif num_controlnet > 1:342            if isinstance(controlnet_conditioning_scale, list):343                if any(isinstance(i, list) for i in controlnet_conditioning_scale):344                    raise ValueError("A single batch of multiple conditionings are supported at the moment.")345            elif (346                isinstance(controlnet_conditioning_scale, list)347                and len(controlnet_conditioning_scale) != num_controlnet348            ):349                raise ValueError(350                    "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"351                    " the same length as the number of controlnets"352                )353        else:354            assert False355 356        if len(control_guidance_start) != len(control_guidance_end):357            raise ValueError(358                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."359            )360 361        if num_controlnet > 1:362            if len(control_guidance_start) != num_controlnet:363                raise ValueError(364                    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}."365                )366 367        for start, end in zip(control_guidance_start, control_guidance_end):368            if start >= end:369                raise ValueError(370                    f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."371                )372            if start < 0.0:373                raise ValueError(f"control guidance start: {start} can't be smaller than 0.")374            if end > 1.0:375                raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")376 377    # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.check_image378    def check_image(self, image, prompt, prompt_embeds):379        image_is_pil = isinstance(image, PIL.Image.Image)380        image_is_tensor = isinstance(image, torch.Tensor)381        image_is_np = isinstance(image, np.ndarray)382        image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)383        image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)384        image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)385 386        if (387            not image_is_pil388            and not image_is_tensor389            and not image_is_np390            and not image_is_pil_list391            and not image_is_tensor_list392            and not image_is_np_list393        ):394            raise TypeError(395                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)}"396            )397 398        if image_is_pil:399            image_batch_size = 1400        else:401            image_batch_size = len(image)402 403        if prompt is not None and isinstance(prompt, str):404            prompt_batch_size = 1405        elif prompt is not None and isinstance(prompt, list):406            prompt_batch_size = len(prompt)407        elif prompt_embeds is not None:408            prompt_batch_size = prompt_embeds.shape[0]409 410        if image_batch_size != 1 and image_batch_size != prompt_batch_size:411            raise ValueError(412                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}"413            )414 415    # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image416    def prepare_control_image(417        self,418        image,419        width,420        height,421        batch_size,422        num_images_per_prompt,423        device,424        dtype,425        do_classifier_free_guidance=False,426        guess_mode=False,427    ):428        image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)429        image_batch_size = image.shape[0]430 431        if image_batch_size == 1:432            repeat_by = batch_size433        else:434            # image batch size is the same as prompt batch size435            repeat_by = num_images_per_prompt436 437        image = image.repeat_interleave(repeat_by, dim=0)438 439        image = image.to(device=device, dtype=dtype)440 441        if do_classifier_free_guidance and not guess_mode:442            image = torch.cat([image] * 2)443 444        return image445 446    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps447    def get_timesteps(self, num_inference_steps, strength, device):448        # get the original timestep using init_timestep449        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)450 451        t_start = max(num_inference_steps - init_timestep, 0)452        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]453 454        return timesteps, num_inference_steps - t_start455 456    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):457        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):458            raise ValueError(459                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"460            )461 462        image = image.to(device=device, dtype=dtype)463 464        batch_size = batch_size * num_images_per_prompt465 466        if image.shape[1] == 4:467            init_latents = image468 469        else:470            _image = image.cpu().detach().numpy()471            init_latents = self.vae_encoder(sample=_image)[0]472            init_latents = torch.from_numpy(init_latents).to(device=device, dtype=dtype)473            init_latents = 0.18215 * init_latents474 475        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:476            # expand init_latents for batch_size477            deprecation_message = (478                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"479                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"480                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"481                " your script to pass as many initial images as text prompts to suppress this warning."482            )483            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)484            additional_image_per_prompt = batch_size // init_latents.shape[0]485            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)486        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:487            raise ValueError(488                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."489            )490        else:491            init_latents = torch.cat([init_latents], dim=0)492 493        shape = init_latents.shape494        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)495 496        # get latents497        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)498        latents = init_latents499 500        return latents501 502    @torch.no_grad()503    @replace_example_docstring(EXAMPLE_DOC_STRING)504    def __call__(505        self,506        num_controlnet: int,507        fp16: bool = True,508        prompt: Union[str, List[str]] = None,509        image: Union[510            torch.Tensor,511            PIL.Image.Image,512            np.ndarray,513            List[torch.Tensor],514            List[PIL.Image.Image],515            List[np.ndarray],516        ] = None,517        control_image: Union[518            torch.Tensor,519            PIL.Image.Image,520            np.ndarray,521            List[torch.Tensor],522            List[PIL.Image.Image],523            List[np.ndarray],524        ] = None,525        height: Optional[int] = None,526        width: Optional[int] = None,527        strength: float = 0.8,528        num_inference_steps: int = 50,529        guidance_scale: float = 7.5,530        negative_prompt: Optional[Union[str, List[str]]] = None,531        num_images_per_prompt: Optional[int] = 1,532        eta: float = 0.0,533        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,534        latents: Optional[torch.Tensor] = None,535        prompt_embeds: Optional[torch.Tensor] = None,536        negative_prompt_embeds: Optional[torch.Tensor] = None,537        output_type: Optional[str] = "pil",538        return_dict: bool = True,539        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,540        callback_steps: int = 1,541        cross_attention_kwargs: Optional[Dict[str, Any]] = None,542        controlnet_conditioning_scale: Union[float, List[float]] = 0.8,543        guess_mode: bool = False,544        control_guidance_start: Union[float, List[float]] = 0.0,545        control_guidance_end: Union[float, List[float]] = 1.0,546    ):547        r"""548        Function invoked when calling the pipeline for generation.549 550        Args:551            prompt (`str` or `List[str]`, *optional*):552                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.553                instead.554            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:555                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):556                The initial image will be used as the starting point for the image generation process. Can also accept557                image latents as `image`, if passing latents directly, it will not be encoded again.558            control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:559                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):560                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If561                the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can562                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If563                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are564                specified in init, images must be passed as a list such that each element of the list can be correctly565                batched for input to a single controlnet.566            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):567                The height in pixels of the generated image.568            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):569                The width in pixels of the generated image.570            num_inference_steps (`int`, *optional*, defaults to 50):571                The number of denoising steps. More denoising steps usually lead to a higher quality image at the572                expense of slower inference.573            guidance_scale (`float`, *optional*, defaults to 7.5):574                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).575                `guidance_scale` is defined as `w` of equation 2. of [Imagen576                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >577                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,578                usually at the expense of lower image quality.579            negative_prompt (`str` or `List[str]`, *optional*):580                The prompt or prompts not to guide the image generation. If not defined, one has to pass581                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is582                less than `1`).583            num_images_per_prompt (`int`, *optional*, defaults to 1):584                The number of images to generate per prompt.585            eta (`float`, *optional*, defaults to 0.0):586                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to587                [`schedulers.DDIMScheduler`], will be ignored for others.588            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):589                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)590                to make generation deterministic.591            latents (`torch.Tensor`, *optional*):592                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image593                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents594                tensor will ge generated by sampling using the supplied random `generator`.595            prompt_embeds (`torch.Tensor`, *optional*):596                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not597                provided, text embeddings will be generated from `prompt` input argument.598            negative_prompt_embeds (`torch.Tensor`, *optional*):599                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt600                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input601                argument.602            output_type (`str`, *optional*, defaults to `"pil"`):603                The output format of the generate image. Choose between604                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.605            return_dict (`bool`, *optional*, defaults to `True`):606                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a607                plain tuple.608            callback (`Callable`, *optional*):609                A function that will be called every `callback_steps` steps during inference. The function will be610                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.611            callback_steps (`int`, *optional*, defaults to 1):612                The frequency at which the `callback` function will be called. If not specified, the callback will be613                called at every step.614            cross_attention_kwargs (`dict`, *optional*):615                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under616                `self.processor` in617                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).618            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):619                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added620                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the621                corresponding scale as a list. Note that by default, we use a smaller conditioning scale for inpainting622                than for [`~StableDiffusionControlNetPipeline.__call__`].623            guess_mode (`bool`, *optional*, defaults to `False`):624                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if625                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.626            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):627                The percentage of total steps at which the controlnet starts applying.628            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):629                The percentage of total steps at which the controlnet stops applying.630 631        Examples:632 633        Returns:634            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:635            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.636            When returning a tuple, the first element is a list with the generated images, and the second element is a637            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"638            (nsfw) content, according to the `safety_checker`.639        """640        if fp16:641            torch_dtype = torch.float16642            np_dtype = np.float16643        else:644            torch_dtype = torch.float32645            np_dtype = np.float32646 647        # align format for control guidance648        if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):649            control_guidance_start = len(control_guidance_end) * [control_guidance_start]650        elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):651            control_guidance_end = len(control_guidance_start) * [control_guidance_end]652        elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):653            mult = num_controlnet654            control_guidance_start, control_guidance_end = (655                mult * [control_guidance_start],656                mult * [control_guidance_end],657            )658 659        # 1. Check inputs. Raise error if not correct660        self.check_inputs(661            num_controlnet,662            prompt,663            control_image,664            callback_steps,665            negative_prompt,666            prompt_embeds,667            negative_prompt_embeds,668            controlnet_conditioning_scale,669            control_guidance_start,670            control_guidance_end,671        )672 673        # 2. Define call parameters674        if prompt is not None and isinstance(prompt, str):675            batch_size = 1676        elif prompt is not None and isinstance(prompt, list):677            batch_size = len(prompt)678        else:679            batch_size = prompt_embeds.shape[0]680 681        device = self._execution_device682        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)683        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`684        # corresponds to doing no classifier free guidance.685        do_classifier_free_guidance = guidance_scale > 1.0686 687        if num_controlnet > 1 and isinstance(controlnet_conditioning_scale, float):688            controlnet_conditioning_scale = [controlnet_conditioning_scale] * num_controlnet689 690        # 3. Encode input prompt691        prompt_embeds = self._encode_prompt(692            prompt,693            num_images_per_prompt,694            do_classifier_free_guidance,695            negative_prompt,696            prompt_embeds=prompt_embeds,697            negative_prompt_embeds=negative_prompt_embeds,698        )699        # 4. Prepare image700        image = self.image_processor.preprocess(image).to(dtype=torch.float32)701 702        # 5. Prepare controlnet_conditioning_image703        if num_controlnet == 1:704            control_image = self.prepare_control_image(705                image=control_image,706                width=width,707                height=height,708                batch_size=batch_size * num_images_per_prompt,709                num_images_per_prompt=num_images_per_prompt,710                device=device,711                dtype=torch_dtype,712                do_classifier_free_guidance=do_classifier_free_guidance,713                guess_mode=guess_mode,714            )715        elif num_controlnet > 1:716            control_images = []717 718            for control_image_ in control_image:719                control_image_ = self.prepare_control_image(720                    image=control_image_,721                    width=width,722                    height=height,723                    batch_size=batch_size * num_images_per_prompt,724                    num_images_per_prompt=num_images_per_prompt,725                    device=device,726                    dtype=torch_dtype,727                    do_classifier_free_guidance=do_classifier_free_guidance,728                    guess_mode=guess_mode,729                )730 731                control_images.append(control_image_)732 733            control_image = control_images734        else:735            assert False736 737        # 5. Prepare timesteps738        self.scheduler.set_timesteps(num_inference_steps, device=device)739        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)740        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)741 742        # 6. Prepare latent variables743        latents = self.prepare_latents(744            image,745            latent_timestep,746            batch_size,747            num_images_per_prompt,748            torch_dtype,749            device,750            generator,751        )752 753        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline754        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)755 756        # 7.1 Create tensor stating which controlnets to keep757        controlnet_keep = []758        for i in range(len(timesteps)):759            keeps = [760                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)761                for s, e in zip(control_guidance_start, control_guidance_end)762            ]763            controlnet_keep.append(keeps[0] if num_controlnet == 1 else keeps)764 765        # 8. Denoising loop766        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order767        with self.progress_bar(total=num_inference_steps) as progress_bar:768            for i, t in enumerate(timesteps):769                # expand the latents if we are doing classifier free guidance770                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents771                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)772 773                if isinstance(controlnet_keep[i], list):774                    cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]775                else:776                    controlnet_cond_scale = controlnet_conditioning_scale777                    if isinstance(controlnet_cond_scale, list):778                        controlnet_cond_scale = controlnet_cond_scale[0]779                    cond_scale = controlnet_cond_scale * controlnet_keep[i]780 781                # predict the noise residual782                _latent_model_input = latent_model_input.cpu().detach().numpy()783                _prompt_embeds = np.array(prompt_embeds, dtype=np_dtype)784                _t = np.array([t.cpu().detach().numpy()], dtype=np_dtype)785 786                if num_controlnet == 1:787                    control_images = np.array([control_image], dtype=np_dtype)788                else:789                    control_images = []790                    for _control_img in control_image:791                        _control_img = _control_img.cpu().detach().numpy()792                        control_images.append(_control_img)793                    control_images = np.array(control_images, dtype=np_dtype)794 795                control_scales = np.array(cond_scale, dtype=np_dtype)796                control_scales = np.resize(control_scales, (num_controlnet, 1))797 798                noise_pred = self.unet(799                    sample=_latent_model_input,800                    timestep=_t,801                    encoder_hidden_states=_prompt_embeds,802                    controlnet_conds=control_images,803                    conditioning_scales=control_scales,804                )[0]805                noise_pred = torch.from_numpy(noise_pred).to(device)806 807                # perform guidance808                if do_classifier_free_guidance:809                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)810                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)811 812                # compute the previous noisy sample x_t -> x_t-1813                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]814 815                # call the callback, if provided816                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):817                    progress_bar.update()818                    if callback is not None and i % callback_steps == 0:819                        step_idx = i // getattr(self.scheduler, "order", 1)820                        callback(step_idx, t, latents)821 822        if not output_type == "latent":823            _latents = latents.cpu().detach().numpy() / 0.18215824            _latents = np.array(_latents, dtype=np_dtype)825            image = self.vae_decoder(latent_sample=_latents)[0]826            image = torch.from_numpy(image).to(device, dtype=torch.float32)827            has_nsfw_concept = None828        else:829            image = latents830            has_nsfw_concept = None831 832        if has_nsfw_concept is None:833            do_denormalize = [True] * image.shape[0]834        else:835            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]836 837        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)838 839        if not return_dict:840            return (image, has_nsfw_concept)841 842        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)843 844 845if __name__ == "__main__":846    parser = argparse.ArgumentParser()847 848    parser.add_argument(849        "--sd_model",850        type=str,851        required=True,852        help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",853    )854 855    parser.add_argument(856        "--onnx_model_dir",857        type=str,858        required=True,859        help="Path to the ONNX directory",860    )861 862    parser.add_argument("--qr_img_path", type=str, required=True, help="Path to the qr code image")863 864    args = parser.parse_args()865 866    qr_image = Image.open(args.qr_img_path)867    qr_image = qr_image.resize((512, 512))868 869    # init stable diffusion pipeline870    pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(args.sd_model)871    pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)872 873    provider = ["CUDAExecutionProvider", "CPUExecutionProvider"]874    onnx_pipeline = OnnxStableDiffusionControlNetImg2ImgPipeline(875        vae_encoder=OnnxRuntimeModel.from_pretrained(876            os.path.join(args.onnx_model_dir, "vae_encoder"), provider=provider877        ),878        vae_decoder=OnnxRuntimeModel.from_pretrained(879            os.path.join(args.onnx_model_dir, "vae_decoder"), provider=provider880        ),881        text_encoder=OnnxRuntimeModel.from_pretrained(882            os.path.join(args.onnx_model_dir, "text_encoder"), provider=provider883        ),884        tokenizer=pipeline.tokenizer,885        unet=OnnxRuntimeModel.from_pretrained(os.path.join(args.onnx_model_dir, "unet"), provider=provider),886        scheduler=pipeline.scheduler,887    )888    onnx_pipeline = onnx_pipeline.to("cuda")889 890    prompt = "a cute cat fly to the moon"891    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"892 893    for i in range(10):894        start_time = time.time()895        image = onnx_pipeline(896            num_controlnet=2,897            prompt=prompt,898            negative_prompt=negative_prompt,899            image=qr_image,900            control_image=[qr_image, qr_image],901            width=512,902            height=512,903            strength=0.75,904            num_inference_steps=20,905            num_images_per_prompt=1,906            controlnet_conditioning_scale=[0.8, 0.8],907            control_guidance_start=[0.3, 0.3],908            control_guidance_end=[0.9, 0.9],909        ).images[0]910        print(time.time() - start_time)911        image.save("output_qr_code.png")912