diffusers/community-pipelines-mirror

1import inspect2import re3from typing import Callable, List, Optional, Union4 5import numpy as np6import PIL7import torch8from packaging import version9from transformers import CLIPImageProcessor, CLIPTokenizer10 11import diffusers12from diffusers import OnnxRuntimeModel, OnnxStableDiffusionPipeline, SchedulerMixin13from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput14from diffusers.utils import logging15 16 17try:18    from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE19except ImportError:20    ORT_TO_NP_TYPE = {21        "tensor(bool)": np.bool_,22        "tensor(int8)": np.int8,23        "tensor(uint8)": np.uint8,24        "tensor(int16)": np.int16,25        "tensor(uint16)": np.uint16,26        "tensor(int32)": np.int32,27        "tensor(uint32)": np.uint32,28        "tensor(int64)": np.int64,29        "tensor(uint64)": np.uint64,30        "tensor(float16)": np.float16,31        "tensor(float)": np.float32,32        "tensor(double)": np.float64,33    }34 35try:36    from diffusers.utils import PIL_INTERPOLATION37except ImportError:38    if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):39        PIL_INTERPOLATION = {40            "linear": PIL.Image.Resampling.BILINEAR,41            "bilinear": PIL.Image.Resampling.BILINEAR,42            "bicubic": PIL.Image.Resampling.BICUBIC,43            "lanczos": PIL.Image.Resampling.LANCZOS,44            "nearest": PIL.Image.Resampling.NEAREST,45        }46    else:47        PIL_INTERPOLATION = {48            "linear": PIL.Image.LINEAR,49            "bilinear": PIL.Image.BILINEAR,50            "bicubic": PIL.Image.BICUBIC,51            "lanczos": PIL.Image.LANCZOS,52            "nearest": PIL.Image.NEAREST,53        }54# ------------------------------------------------------------------------------55 56logger = logging.get_logger(__name__)  # pylint: disable=invalid-name57 58re_attention = re.compile(59    r"""60\\\(|61\\\)|62\\\[|63\\]|64\\\\|65\\|66\(|67\[|68:([+-]?[.\d]+)\)|69\)|70]|71[^\\()\[\]:]+|72:73""",74    re.X,75)76 77 78def parse_prompt_attention(text):79    """80    Parses a string with attention tokens and returns a list of pairs: text and its associated weight.81    Accepted tokens are:82      (abc) - increases attention to abc by a multiplier of 1.183      (abc:3.12) - increases attention to abc by a multiplier of 3.1284      [abc] - decreases attention to abc by a multiplier of 1.185      \( - literal character '('86      \[ - literal character '['87      \) - literal character ')'88      \] - literal character ']'89      \\ - literal character '\'90      anything else - just text91    >>> parse_prompt_attention('normal text')92    [['normal text', 1.0]]93    >>> parse_prompt_attention('an (important) word')94    [['an ', 1.0], ['important', 1.1], [' word', 1.0]]95    >>> parse_prompt_attention('(unbalanced')96    [['unbalanced', 1.1]]97    >>> parse_prompt_attention('\(literal\]')98    [['(literal]', 1.0]]99    >>> parse_prompt_attention('(unnecessary)(parens)')100    [['unnecessaryparens', 1.1]]101    >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')102    [['a ', 1.0],103     ['house', 1.5730000000000004],104     [' ', 1.1],105     ['on', 1.0],106     [' a ', 1.1],107     ['hill', 0.55],108     [', sun, ', 1.1],109     ['sky', 1.4641000000000006],110     ['.', 1.1]]111    """112 113    res = []114    round_brackets = []115    square_brackets = []116 117    round_bracket_multiplier = 1.1118    square_bracket_multiplier = 1 / 1.1119 120    def multiply_range(start_position, multiplier):121        for p in range(start_position, len(res)):122            res[p][1] *= multiplier123 124    for m in re_attention.finditer(text):125        text = m.group(0)126        weight = m.group(1)127 128        if text.startswith("\\"):129            res.append([text[1:], 1.0])130        elif text == "(":131            round_brackets.append(len(res))132        elif text == "[":133            square_brackets.append(len(res))134        elif weight is not None and len(round_brackets) > 0:135            multiply_range(round_brackets.pop(), float(weight))136        elif text == ")" and len(round_brackets) > 0:137            multiply_range(round_brackets.pop(), round_bracket_multiplier)138        elif text == "]" and len(square_brackets) > 0:139            multiply_range(square_brackets.pop(), square_bracket_multiplier)140        else:141            res.append([text, 1.0])142 143    for pos in round_brackets:144        multiply_range(pos, round_bracket_multiplier)145 146    for pos in square_brackets:147        multiply_range(pos, square_bracket_multiplier)148 149    if len(res) == 0:150        res = [["", 1.0]]151 152    # merge runs of identical weights153    i = 0154    while i + 1 < len(res):155        if res[i][1] == res[i + 1][1]:156            res[i][0] += res[i + 1][0]157            res.pop(i + 1)158        else:159            i += 1160 161    return res162 163 164def get_prompts_with_weights(pipe, prompt: List[str], max_length: int):165    r"""166    Tokenize a list of prompts and return its tokens with weights of each token.167 168    No padding, starting or ending token is included.169    """170    tokens = []171    weights = []172    truncated = False173    for text in prompt:174        texts_and_weights = parse_prompt_attention(text)175        text_token = []176        text_weight = []177        for word, weight in texts_and_weights:178            # tokenize and discard the starting and the ending token179            token = pipe.tokenizer(word, return_tensors="np").input_ids[0, 1:-1]180            text_token += list(token)181            # copy the weight by length of token182            text_weight += [weight] * len(token)183            # stop if the text is too long (longer than truncation limit)184            if len(text_token) > max_length:185                truncated = True186                break187        # truncate188        if len(text_token) > max_length:189            truncated = True190            text_token = text_token[:max_length]191            text_weight = text_weight[:max_length]192        tokens.append(text_token)193        weights.append(text_weight)194    if truncated:195        logger.warning("Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples")196    return tokens, weights197 198 199def pad_tokens_and_weights(tokens, weights, max_length, bos, eos, pad, no_boseos_middle=True, chunk_length=77):200    r"""201    Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length.202    """203    max_embeddings_multiples = (max_length - 2) // (chunk_length - 2)204    weights_length = max_length if no_boseos_middle else max_embeddings_multiples * chunk_length205    for i in range(len(tokens)):206        tokens[i] = [bos] + tokens[i] + [pad] * (max_length - 1 - len(tokens[i]) - 1) + [eos]207        if no_boseos_middle:208            weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i]))209        else:210            w = []211            if len(weights[i]) == 0:212                w = [1.0] * weights_length213            else:214                for j in range(max_embeddings_multiples):215                    w.append(1.0)  # weight for starting token in this chunk216                    w += weights[i][j * (chunk_length - 2) : min(len(weights[i]), (j + 1) * (chunk_length - 2))]217                    w.append(1.0)  # weight for ending token in this chunk218                w += [1.0] * (weights_length - len(w))219            weights[i] = w[:]220 221    return tokens, weights222 223 224def get_unweighted_text_embeddings(225    pipe,226    text_input: np.array,227    chunk_length: int,228    no_boseos_middle: Optional[bool] = True,229):230    """231    When the length of tokens is a multiple of the capacity of the text encoder,232    it should be split into chunks and sent to the text encoder individually.233    """234    max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2)235    if max_embeddings_multiples > 1:236        text_embeddings = []237        for i in range(max_embeddings_multiples):238            # extract the i-th chunk239            text_input_chunk = text_input[:, i * (chunk_length - 2) : (i + 1) * (chunk_length - 2) + 2].copy()240 241            # cover the head and the tail by the starting and the ending tokens242            text_input_chunk[:, 0] = text_input[0, 0]243            text_input_chunk[:, -1] = text_input[0, -1]244 245            text_embedding = pipe.text_encoder(input_ids=text_input_chunk)[0]246 247            if no_boseos_middle:248                if i == 0:249                    # discard the ending token250                    text_embedding = text_embedding[:, :-1]251                elif i == max_embeddings_multiples - 1:252                    # discard the starting token253                    text_embedding = text_embedding[:, 1:]254                else:255                    # discard both starting and ending tokens256                    text_embedding = text_embedding[:, 1:-1]257 258            text_embeddings.append(text_embedding)259        text_embeddings = np.concatenate(text_embeddings, axis=1)260    else:261        text_embeddings = pipe.text_encoder(input_ids=text_input)[0]262    return text_embeddings263 264 265def get_weighted_text_embeddings(266    pipe,267    prompt: Union[str, List[str]],268    uncond_prompt: Optional[Union[str, List[str]]] = None,269    max_embeddings_multiples: Optional[int] = 4,270    no_boseos_middle: Optional[bool] = False,271    skip_parsing: Optional[bool] = False,272    skip_weighting: Optional[bool] = False,273    **kwargs,274):275    r"""276    Prompts can be assigned with local weights using brackets. For example,277    prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful',278    and the embedding tokens corresponding to the words get multiplied by a constant, 1.1.279 280    Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean.281 282    Args:283        pipe (`OnnxStableDiffusionPipeline`):284            Pipe to provide access to the tokenizer and the text encoder.285        prompt (`str` or `List[str]`):286            The prompt or prompts to guide the image generation.287        uncond_prompt (`str` or `List[str]`):288            The unconditional prompt or prompts for guide the image generation. If unconditional prompt289            is provided, the embeddings of prompt and uncond_prompt are concatenated.290        max_embeddings_multiples (`int`, *optional*, defaults to `1`):291            The max multiple length of prompt embeddings compared to the max output length of text encoder.292        no_boseos_middle (`bool`, *optional*, defaults to `False`):293            If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and294            ending token in each of the chunk in the middle.295        skip_parsing (`bool`, *optional*, defaults to `False`):296            Skip the parsing of brackets.297        skip_weighting (`bool`, *optional*, defaults to `False`):298            Skip the weighting. When the parsing is skipped, it is forced True.299    """300    max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2301    if isinstance(prompt, str):302        prompt = [prompt]303 304    if not skip_parsing:305        prompt_tokens, prompt_weights = get_prompts_with_weights(pipe, prompt, max_length - 2)306        if uncond_prompt is not None:307            if isinstance(uncond_prompt, str):308                uncond_prompt = [uncond_prompt]309            uncond_tokens, uncond_weights = get_prompts_with_weights(pipe, uncond_prompt, max_length - 2)310    else:311        prompt_tokens = [312            token[1:-1]313            for token in pipe.tokenizer(prompt, max_length=max_length, truncation=True, return_tensors="np").input_ids314        ]315        prompt_weights = [[1.0] * len(token) for token in prompt_tokens]316        if uncond_prompt is not None:317            if isinstance(uncond_prompt, str):318                uncond_prompt = [uncond_prompt]319            uncond_tokens = [320                token[1:-1]321                for token in pipe.tokenizer(322                    uncond_prompt,323                    max_length=max_length,324                    truncation=True,325                    return_tensors="np",326                ).input_ids327            ]328            uncond_weights = [[1.0] * len(token) for token in uncond_tokens]329 330    # round up the longest length of tokens to a multiple of (model_max_length - 2)331    max_length = max([len(token) for token in prompt_tokens])332    if uncond_prompt is not None:333        max_length = max(max_length, max([len(token) for token in uncond_tokens]))334 335    max_embeddings_multiples = min(336        max_embeddings_multiples,337        (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1,338    )339    max_embeddings_multiples = max(1, max_embeddings_multiples)340    max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2341 342    # pad the length of tokens and weights343    bos = pipe.tokenizer.bos_token_id344    eos = pipe.tokenizer.eos_token_id345    pad = getattr(pipe.tokenizer, "pad_token_id", eos)346    prompt_tokens, prompt_weights = pad_tokens_and_weights(347        prompt_tokens,348        prompt_weights,349        max_length,350        bos,351        eos,352        pad,353        no_boseos_middle=no_boseos_middle,354        chunk_length=pipe.tokenizer.model_max_length,355    )356    prompt_tokens = np.array(prompt_tokens, dtype=np.int32)357    if uncond_prompt is not None:358        uncond_tokens, uncond_weights = pad_tokens_and_weights(359            uncond_tokens,360            uncond_weights,361            max_length,362            bos,363            eos,364            pad,365            no_boseos_middle=no_boseos_middle,366            chunk_length=pipe.tokenizer.model_max_length,367        )368        uncond_tokens = np.array(uncond_tokens, dtype=np.int32)369 370    # get the embeddings371    text_embeddings = get_unweighted_text_embeddings(372        pipe,373        prompt_tokens,374        pipe.tokenizer.model_max_length,375        no_boseos_middle=no_boseos_middle,376    )377    prompt_weights = np.array(prompt_weights, dtype=text_embeddings.dtype)378    if uncond_prompt is not None:379        uncond_embeddings = get_unweighted_text_embeddings(380            pipe,381            uncond_tokens,382            pipe.tokenizer.model_max_length,383            no_boseos_middle=no_boseos_middle,384        )385        uncond_weights = np.array(uncond_weights, dtype=uncond_embeddings.dtype)386 387    # assign weights to the prompts and normalize in the sense of mean388    # TODO: should we normalize by chunk or in a whole (current implementation)?389    if (not skip_parsing) and (not skip_weighting):390        previous_mean = text_embeddings.mean(axis=(-2, -1))391        text_embeddings *= prompt_weights[:, :, None]392        text_embeddings *= (previous_mean / text_embeddings.mean(axis=(-2, -1)))[:, None, None]393        if uncond_prompt is not None:394            previous_mean = uncond_embeddings.mean(axis=(-2, -1))395            uncond_embeddings *= uncond_weights[:, :, None]396            uncond_embeddings *= (previous_mean / uncond_embeddings.mean(axis=(-2, -1)))[:, None, None]397 398    # For classifier free guidance, we need to do two forward passes.399    # Here we concatenate the unconditional and text embeddings into a single batch400    # to avoid doing two forward passes401    if uncond_prompt is not None:402        return text_embeddings, uncond_embeddings403 404    return text_embeddings405 406 407def preprocess_image(image):408    w, h = image.size409    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32410    image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"])411    image = np.array(image).astype(np.float32) / 255.0412    image = image[None].transpose(0, 3, 1, 2)413    return 2.0 * image - 1.0414 415 416def preprocess_mask(mask, scale_factor=8):417    mask = mask.convert("L")418    w, h = mask.size419    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32420    mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL_INTERPOLATION["nearest"])421    mask = np.array(mask).astype(np.float32) / 255.0422    mask = np.tile(mask, (4, 1, 1))423    mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?424    mask = 1 - mask  # repaint white, keep black425    return mask426 427 428class OnnxStableDiffusionLongPromptWeightingPipeline(OnnxStableDiffusionPipeline):429    r"""430    Pipeline for text-to-image generation using Stable Diffusion without tokens length limit, and support parsing431    weighting in prompt.432 433    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the434    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)435    """436    if version.parse(version.parse(diffusers.__version__).base_version) >= version.parse("0.9.0"):437 438        def __init__(439            self,440            vae_encoder: OnnxRuntimeModel,441            vae_decoder: OnnxRuntimeModel,442            text_encoder: OnnxRuntimeModel,443            tokenizer: CLIPTokenizer,444            unet: OnnxRuntimeModel,445            scheduler: SchedulerMixin,446            safety_checker: OnnxRuntimeModel,447            feature_extractor: CLIPImageProcessor,448            requires_safety_checker: bool = True,449        ):450            super().__init__(451                vae_encoder=vae_encoder,452                vae_decoder=vae_decoder,453                text_encoder=text_encoder,454                tokenizer=tokenizer,455                unet=unet,456                scheduler=scheduler,457                safety_checker=safety_checker,458                feature_extractor=feature_extractor,459                requires_safety_checker=requires_safety_checker,460            )461            self.__init__additional__()462 463    else:464 465        def __init__(466            self,467            vae_encoder: OnnxRuntimeModel,468            vae_decoder: OnnxRuntimeModel,469            text_encoder: OnnxRuntimeModel,470            tokenizer: CLIPTokenizer,471            unet: OnnxRuntimeModel,472            scheduler: SchedulerMixin,473            safety_checker: OnnxRuntimeModel,474            feature_extractor: CLIPImageProcessor,475        ):476            super().__init__(477                vae_encoder=vae_encoder,478                vae_decoder=vae_decoder,479                text_encoder=text_encoder,480                tokenizer=tokenizer,481                unet=unet,482                scheduler=scheduler,483                safety_checker=safety_checker,484                feature_extractor=feature_extractor,485            )486            self.__init__additional__()487 488    def __init__additional__(self):489        self.unet.config.in_channels = 4490        self.vae_scale_factor = 8491 492    def _encode_prompt(493        self,494        prompt,495        num_images_per_prompt,496        do_classifier_free_guidance,497        negative_prompt,498        max_embeddings_multiples,499    ):500        r"""501        Encodes the prompt into text encoder hidden states.502 503        Args:504            prompt (`str` or `list(int)`):505                prompt to be encoded506            num_images_per_prompt (`int`):507                number of images that should be generated per prompt508            do_classifier_free_guidance (`bool`):509                whether to use classifier free guidance or not510            negative_prompt (`str` or `List[str]`):511                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored512                if `guidance_scale` is less than `1`).513            max_embeddings_multiples (`int`, *optional*, defaults to `3`):514                The max multiple length of prompt embeddings compared to the max output length of text encoder.515        """516        batch_size = len(prompt) if isinstance(prompt, list) else 1517 518        if negative_prompt is None:519            negative_prompt = [""] * batch_size520        elif isinstance(negative_prompt, str):521            negative_prompt = [negative_prompt] * batch_size522        if batch_size != len(negative_prompt):523            raise ValueError(524                f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"525                f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"526                " the batch size of `prompt`."527            )528 529        text_embeddings, uncond_embeddings = get_weighted_text_embeddings(530            pipe=self,531            prompt=prompt,532            uncond_prompt=negative_prompt if do_classifier_free_guidance else None,533            max_embeddings_multiples=max_embeddings_multiples,534        )535 536        text_embeddings = text_embeddings.repeat(num_images_per_prompt, 0)537        if do_classifier_free_guidance:538            uncond_embeddings = uncond_embeddings.repeat(num_images_per_prompt, 0)539            text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])540 541        return text_embeddings542 543    def check_inputs(self, prompt, height, width, strength, callback_steps):544        if not isinstance(prompt, str) and not isinstance(prompt, list):545            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")546 547        if strength < 0 or strength > 1:548            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")549 550        if height % 8 != 0 or width % 8 != 0:551            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")552 553        if (callback_steps is None) or (554            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)555        ):556            raise ValueError(557                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"558                f" {type(callback_steps)}."559            )560 561    def get_timesteps(self, num_inference_steps, strength, is_text2img):562        if is_text2img:563            return self.scheduler.timesteps, num_inference_steps564        else:565            # get the original timestep using init_timestep566            offset = self.scheduler.config.get("steps_offset", 0)567            init_timestep = int(num_inference_steps * strength) + offset568            init_timestep = min(init_timestep, num_inference_steps)569 570            t_start = max(num_inference_steps - init_timestep + offset, 0)571            timesteps = self.scheduler.timesteps[t_start:]572            return timesteps, num_inference_steps - t_start573 574    def run_safety_checker(self, image):575        if self.safety_checker is not None:576            safety_checker_input = self.feature_extractor(577                self.numpy_to_pil(image), return_tensors="np"578            ).pixel_values.astype(image.dtype)579            # There will throw an error if use safety_checker directly and batchsize>1580            images, has_nsfw_concept = [], []581            for i in range(image.shape[0]):582                image_i, has_nsfw_concept_i = self.safety_checker(583                    clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1]584                )585                images.append(image_i)586                has_nsfw_concept.append(has_nsfw_concept_i[0])587            image = np.concatenate(images)588        else:589            has_nsfw_concept = None590        return image, has_nsfw_concept591 592    def decode_latents(self, latents):593        latents = 1 / 0.18215 * latents594        # image = self.vae_decoder(latent_sample=latents)[0]595        # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1596        image = np.concatenate(597            [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]598        )599        image = np.clip(image / 2 + 0.5, 0, 1)600        image = image.transpose((0, 2, 3, 1))601        return image602 603    def prepare_extra_step_kwargs(self, generator, eta):604        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature605        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.606        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502607        # and should be between [0, 1]608 609        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())610        extra_step_kwargs = {}611        if accepts_eta:612            extra_step_kwargs["eta"] = eta613 614        # check if the scheduler accepts generator615        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())616        if accepts_generator:617            extra_step_kwargs["generator"] = generator618        return extra_step_kwargs619 620    def prepare_latents(self, image, timestep, batch_size, height, width, dtype, generator, latents=None):621        if image is None:622            shape = (623                batch_size,624                self.unet.config.in_channels,625                height // self.vae_scale_factor,626                width // self.vae_scale_factor,627            )628 629            if latents is None:630                latents = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)631            else:632                if latents.shape != shape:633                    raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")634 635            # scale the initial noise by the standard deviation required by the scheduler636            latents = (torch.from_numpy(latents) * self.scheduler.init_noise_sigma).numpy()637            return latents, None, None638        else:639            init_latents = self.vae_encoder(sample=image)[0]640            init_latents = 0.18215 * init_latents641            init_latents = np.concatenate([init_latents] * batch_size, axis=0)642            init_latents_orig = init_latents643            shape = init_latents.shape644 645            # add noise to latents using the timesteps646            noise = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)647            latents = self.scheduler.add_noise(648                torch.from_numpy(init_latents), torch.from_numpy(noise), timestep649            ).numpy()650            return latents, init_latents_orig, noise651 652    @torch.no_grad()653    def __call__(654        self,655        prompt: Union[str, List[str]],656        negative_prompt: Optional[Union[str, List[str]]] = None,657        image: Union[np.ndarray, PIL.Image.Image] = None,658        mask_image: Union[np.ndarray, PIL.Image.Image] = None,659        height: int = 512,660        width: int = 512,661        num_inference_steps: int = 50,662        guidance_scale: float = 7.5,663        strength: float = 0.8,664        num_images_per_prompt: Optional[int] = 1,665        eta: float = 0.0,666        generator: Optional[torch.Generator] = None,667        latents: Optional[np.ndarray] = None,668        max_embeddings_multiples: Optional[int] = 3,669        output_type: Optional[str] = "pil",670        return_dict: bool = True,671        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,672        is_cancelled_callback: Optional[Callable[[], bool]] = None,673        callback_steps: int = 1,674        **kwargs,675    ):676        r"""677        Function invoked when calling the pipeline for generation.678 679        Args:680            prompt (`str` or `List[str]`):681                The prompt or prompts to guide the image generation.682            negative_prompt (`str` or `List[str]`, *optional*):683                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored684                if `guidance_scale` is less than `1`).685            image (`np.ndarray` or `PIL.Image.Image`):686                `Image`, or tensor representing an image batch, that will be used as the starting point for the687                process.688            mask_image (`np.ndarray` or `PIL.Image.Image`):689                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be690                replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a691                PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should692                contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.693            height (`int`, *optional*, defaults to 512):694                The height in pixels of the generated image.695            width (`int`, *optional*, defaults to 512):696                The width in pixels of the generated image.697            num_inference_steps (`int`, *optional*, defaults to 50):698                The number of denoising steps. More denoising steps usually lead to a higher quality image at the699                expense of slower inference.700            guidance_scale (`float`, *optional*, defaults to 7.5):701                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).702                `guidance_scale` is defined as `w` of equation 2. of [Imagen703                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >704                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,705                usually at the expense of lower image quality.706            strength (`float`, *optional*, defaults to 0.8):707                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.708                `image` will be used as a starting point, adding more noise to it the larger the `strength`. The709                number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added710                noise will be maximum and the denoising process will run for the full number of iterations specified in711                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.712            num_images_per_prompt (`int`, *optional*, defaults to 1):713                The number of images to generate per prompt.714            eta (`float`, *optional*, defaults to 0.0):715                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to716                [`schedulers.DDIMScheduler`], will be ignored for others.717            generator (`torch.Generator`, *optional*):718                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation719                deterministic.720            latents (`np.ndarray`, *optional*):721                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image722                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents723                tensor will ge generated by sampling using the supplied random `generator`.724            max_embeddings_multiples (`int`, *optional*, defaults to `3`):725                The max multiple length of prompt embeddings compared to the max output length of text encoder.726            output_type (`str`, *optional*, defaults to `"pil"`):727                The output format of the generate image. Choose between728                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.729            return_dict (`bool`, *optional*, defaults to `True`):730                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a731                plain tuple.732            callback (`Callable`, *optional*):733                A function that will be called every `callback_steps` steps during inference. The function will be734                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.735            is_cancelled_callback (`Callable`, *optional*):736                A function that will be called every `callback_steps` steps during inference. If the function returns737                `True`, the inference will be cancelled.738            callback_steps (`int`, *optional*, defaults to 1):739                The frequency at which the `callback` function will be called. If not specified, the callback will be740                called at every step.741 742        Returns:743            `None` if cancelled by `is_cancelled_callback`,744            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:745            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.746            When returning a tuple, the first element is a list with the generated images, and the second element is a747            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"748            (nsfw) content, according to the `safety_checker`.749        """750        # 0. Default height and width to unet751        height = height or self.unet.config.sample_size * self.vae_scale_factor752        width = width or self.unet.config.sample_size * self.vae_scale_factor753 754        # 1. Check inputs. Raise error if not correct755        self.check_inputs(prompt, height, width, strength, callback_steps)756 757        # 2. Define call parameters758        batch_size = 1 if isinstance(prompt, str) else len(prompt)759        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)760        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`761        # corresponds to doing no classifier free guidance.762        do_classifier_free_guidance = guidance_scale > 1.0763 764        # 3. Encode input prompt765        text_embeddings = self._encode_prompt(766            prompt,767            num_images_per_prompt,768            do_classifier_free_guidance,769            negative_prompt,770            max_embeddings_multiples,771        )772        dtype = text_embeddings.dtype773 774        # 4. Preprocess image and mask775        if isinstance(image, PIL.Image.Image):776            image = preprocess_image(image)777        if image is not None:778            image = image.astype(dtype)779        if isinstance(mask_image, PIL.Image.Image):780            mask_image = preprocess_mask(mask_image, self.vae_scale_factor)781        if mask_image is not None:782            mask = mask_image.astype(dtype)783            mask = np.concatenate([mask] * batch_size * num_images_per_prompt)784        else:785            mask = None786 787        # 5. set timesteps788        self.scheduler.set_timesteps(num_inference_steps)789        timestep_dtype = next(790            (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)"791        )792        timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype]793        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, image is None)794        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)795 796        # 6. Prepare latent variables797        latents, init_latents_orig, noise = self.prepare_latents(798            image,799            latent_timestep,800            batch_size * num_images_per_prompt,801            height,802            width,803            dtype,804            generator,805            latents,806        )807 808        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline809        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)810 811        # 8. Denoising loop812        for i, t in enumerate(self.progress_bar(timesteps)):813            # expand the latents if we are doing classifier free guidance814            latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents815            latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)816            latent_model_input = latent_model_input.numpy()817 818            # predict the noise residual819            noise_pred = self.unet(820                sample=latent_model_input,821                timestep=np.array([t], dtype=timestep_dtype),822                encoder_hidden_states=text_embeddings,823            )824            noise_pred = noise_pred[0]825 826            # perform guidance827            if do_classifier_free_guidance:828                noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)829                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)830 831            # compute the previous noisy sample x_t -> x_t-1832            scheduler_output = self.scheduler.step(833                torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs834            )835            latents = scheduler_output.prev_sample.numpy()836 837            if mask is not None:838                # masking839                init_latents_proper = self.scheduler.add_noise(840                    torch.from_numpy(init_latents_orig),841                    torch.from_numpy(noise),842                    t,843                ).numpy()844                latents = (init_latents_proper * mask) + (latents * (1 - mask))845 846            # call the callback, if provided847            if i % callback_steps == 0:848                if callback is not None:849                    callback(i, t, latents)850                if is_cancelled_callback is not None and is_cancelled_callback():851                    return None852 853        # 9. Post-processing854        image = self.decode_latents(latents)855 856        # 10. Run safety checker857        image, has_nsfw_concept = self.run_safety_checker(image)858 859        # 11. Convert to PIL860        if output_type == "pil":861            image = self.numpy_to_pil(image)862 863        if not return_dict:864            return image, has_nsfw_concept865 866        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)867 868    def text2img(869        self,870        prompt: Union[str, List[str]],871        negative_prompt: Optional[Union[str, List[str]]] = None,872        height: int = 512,873        width: int = 512,874        num_inference_steps: int = 50,875        guidance_scale: float = 7.5,876        num_images_per_prompt: Optional[int] = 1,877        eta: float = 0.0,878        generator: Optional[torch.Generator] = None,879        latents: Optional[np.ndarray] = None,880        max_embeddings_multiples: Optional[int] = 3,881        output_type: Optional[str] = "pil",882        return_dict: bool = True,883        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,884        callback_steps: int = 1,885        **kwargs,886    ):887        r"""888        Function for text-to-image generation.889        Args:890            prompt (`str` or `List[str]`):891                The prompt or prompts to guide the image generation.892            negative_prompt (`str` or `List[str]`, *optional*):893                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored894                if `guidance_scale` is less than `1`).895            height (`int`, *optional*, defaults to 512):896                The height in pixels of the generated image.897            width (`int`, *optional*, defaults to 512):898                The width in pixels of the generated image.899            num_inference_steps (`int`, *optional*, defaults to 50):900                The number of denoising steps. More denoising steps usually lead to a higher quality image at the901                expense of slower inference.902            guidance_scale (`float`, *optional*, defaults to 7.5):903                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).904                `guidance_scale` is defined as `w` of equation 2. of [Imagen905                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >906                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,907                usually at the expense of lower image quality.908            num_images_per_prompt (`int`, *optional*, defaults to 1):909                The number of images to generate per prompt.910            eta (`float`, *optional*, defaults to 0.0):911                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to912                [`schedulers.DDIMScheduler`], will be ignored for others.913            generator (`torch.Generator`, *optional*):914                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation915                deterministic.916            latents (`np.ndarray`, *optional*):917                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image918                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents919                tensor will ge generated by sampling using the supplied random `generator`.920            max_embeddings_multiples (`int`, *optional*, defaults to `3`):921                The max multiple length of prompt embeddings compared to the max output length of text encoder.922            output_type (`str`, *optional*, defaults to `"pil"`):923                The output format of the generate image. Choose between924                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.925            return_dict (`bool`, *optional*, defaults to `True`):926                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a927                plain tuple.928            callback (`Callable`, *optional*):929                A function that will be called every `callback_steps` steps during inference. The function will be930                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.931            callback_steps (`int`, *optional*, defaults to 1):932                The frequency at which the `callback` function will be called. If not specified, the callback will be933                called at every step.934        Returns:935            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:936            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.937            When returning a tuple, the first element is a list with the generated images, and the second element is a938            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"939            (nsfw) content, according to the `safety_checker`.940        """941        return self.__call__(942            prompt=prompt,943            negative_prompt=negative_prompt,944            height=height,945            width=width,946            num_inference_steps=num_inference_steps,947            guidance_scale=guidance_scale,948            num_images_per_prompt=num_images_per_prompt,949            eta=eta,950            generator=generator,951            latents=latents,952            max_embeddings_multiples=max_embeddings_multiples,953            output_type=output_type,954            return_dict=return_dict,955            callback=callback,956            callback_steps=callback_steps,957            **kwargs,958        )959 960    def img2img(961        self,962        image: Union[np.ndarray, PIL.Image.Image],963        prompt: Union[str, List[str]],964        negative_prompt: Optional[Union[str, List[str]]] = None,965        strength: float = 0.8,966        num_inference_steps: Optional[int] = 50,967        guidance_scale: Optional[float] = 7.5,968        num_images_per_prompt: Optional[int] = 1,969        eta: Optional[float] = 0.0,970        generator: Optional[torch.Generator] = None,971        max_embeddings_multiples: Optional[int] = 3,972        output_type: Optional[str] = "pil",973        return_dict: bool = True,974        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,975        callback_steps: int = 1,976        **kwargs,977    ):978        r"""979        Function for image-to-image generation.980        Args:981            image (`np.ndarray` or `PIL.Image.Image`):982                `Image`, or ndarray representing an image batch, that will be used as the starting point for the983                process.984            prompt (`str` or `List[str]`):985                The prompt or prompts to guide the image generation.986            negative_prompt (`str` or `List[str]`, *optional*):987                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored988                if `guidance_scale` is less than `1`).989            strength (`float`, *optional*, defaults to 0.8):990                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.991                `image` will be used as a starting point, adding more noise to it the larger the `strength`. The992                number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added993                noise will be maximum and the denoising process will run for the full number of iterations specified in994                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.995            num_inference_steps (`int`, *optional*, defaults to 50):996                The number of denoising steps. More denoising steps usually lead to a higher quality image at the997                expense of slower inference. This parameter will be modulated by `strength`.998            guidance_scale (`float`, *optional*, defaults to 7.5):999                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1000                `guidance_scale` is defined as `w` of equation 2. of [Imagen1001                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1002                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1003                usually at the expense of lower image quality.1004            num_images_per_prompt (`int`, *optional*, defaults to 1):1005                The number of images to generate per prompt.1006            eta (`float`, *optional*, defaults to 0.0):1007                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to1008                [`schedulers.DDIMScheduler`], will be ignored for others.1009            generator (`torch.Generator`, *optional*):1010                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation1011                deterministic.1012            max_embeddings_multiples (`int`, *optional*, defaults to `3`):1013                The max multiple length of prompt embeddings compared to the max output length of text encoder.1014            output_type (`str`, *optional*, defaults to `"pil"`):1015                The output format of the generate image. Choose between1016                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.1017            return_dict (`bool`, *optional*, defaults to `True`):1018                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a1019                plain tuple.1020            callback (`Callable`, *optional*):1021                A function that will be called every `callback_steps` steps during inference. The function will be1022                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.1023            callback_steps (`int`, *optional*, defaults to 1):1024                The frequency at which the `callback` function will be called. If not specified, the callback will be1025                called at every step.1026        Returns:1027            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:1028            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.1029            When returning a tuple, the first element is a list with the generated images, and the second element is a1030            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"1031            (nsfw) content, according to the `safety_checker`.1032        """1033        return self.__call__(1034            prompt=prompt,1035            negative_prompt=negative_prompt,1036            image=image,1037            num_inference_steps=num_inference_steps,1038            guidance_scale=guidance_scale,1039            strength=strength,1040            num_images_per_prompt=num_images_per_prompt,1041            eta=eta,1042            generator=generator,1043            max_embeddings_multiples=max_embeddings_multiples,1044            output_type=output_type,1045            return_dict=return_dict,1046            callback=callback,1047            callback_steps=callback_steps,1048            **kwargs,1049        )1050 1051    def inpaint(1052        self,1053        image: Union[np.ndarray, PIL.Image.Image],1054        mask_image: Union[np.ndarray, PIL.Image.Image],1055        prompt: Union[str, List[str]],1056        negative_prompt: Optional[Union[str, List[str]]] = None,1057        strength: float = 0.8,1058        num_inference_steps: Optional[int] = 50,1059        guidance_scale: Optional[float] = 7.5,1060        num_images_per_prompt: Optional[int] = 1,1061        eta: Optional[float] = 0.0,1062        generator: Optional[torch.Generator] = None,1063        max_embeddings_multiples: Optional[int] = 3,1064        output_type: Optional[str] = "pil",1065        return_dict: bool = True,1066        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,1067        callback_steps: int = 1,1068        **kwargs,1069    ):1070        r"""1071        Function for inpaint.1072        Args:1073            image (`np.ndarray` or `PIL.Image.Image`):1074                `Image`, or tensor representing an image batch, that will be used as the starting point for the1075                process. This is the image whose masked region will be inpainted.1076            mask_image (`np.ndarray` or `PIL.Image.Image`):1077                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be1078                replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a1079                PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should1080                contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.1081            prompt (`str` or `List[str]`):1082                The prompt or prompts to guide the image generation.1083            negative_prompt (`str` or `List[str]`, *optional*):1084                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored1085                if `guidance_scale` is less than `1`).1086            strength (`float`, *optional*, defaults to 0.8):1087                Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength`1088                is 1, the denoising process will be run on the masked area for the full number of iterations specified1089                in `num_inference_steps`. `image` will be used as a reference for the masked area, adding more1090                noise to that region the larger the `strength`. If `strength` is 0, no inpainting will occur.1091            num_inference_steps (`int`, *optional*, defaults to 50):1092                The reference number of denoising steps. More denoising steps usually lead to a higher quality image at1093                the expense of slower inference. This parameter will be modulated by `strength`, as explained above.1094            guidance_scale (`float`, *optional*, defaults to 7.5):1095                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1096                `guidance_scale` is defined as `w` of equation 2. of [Imagen1097                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1098                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1099                usually at the expense of lower image quality.1100            num_images_per_prompt (`int`, *optional*, defaults to 1):1101                The number of images to generate per prompt.1102            eta (`float`, *optional*, defaults to 0.0):1103                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to1104                [`schedulers.DDIMScheduler`], will be ignored for others.1105            generator (`torch.Generator`, *optional*):1106                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation1107                deterministic.1108            max_embeddings_multiples (`int`, *optional*, defaults to `3`):1109                The max multiple length of prompt embeddings compared to the max output length of text encoder.1110            output_type (`str`, *optional*, defaults to `"pil"`):1111                The output format of the generate image. Choose between1112                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.1113            return_dict (`bool`, *optional*, defaults to `True`):1114                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a1115                plain tuple.1116            callback (`Callable`, *optional*):1117                A function that will be called every `callback_steps` steps during inference. The function will be1118                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.1119            callback_steps (`int`, *optional*, defaults to 1):1120                The frequency at which the `callback` function will be called. If not specified, the callback will be1121                called at every step.1122        Returns:1123            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:1124            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.1125            When returning a tuple, the first element is a list with the generated images, and the second element is a1126            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"1127            (nsfw) content, according to the `safety_checker`.1128        """1129        return self.__call__(1130            prompt=prompt,1131            negative_prompt=negative_prompt,1132            image=image,1133            mask_image=mask_image,1134            num_inference_steps=num_inference_steps,1135            guidance_scale=guidance_scale,1136            strength=strength,1137            num_images_per_prompt=num_images_per_prompt,1138            eta=eta,1139            generator=generator,1140            max_embeddings_multiples=max_embeddings_multiples,1141            output_type=output_type,1142            return_dict=return_dict,1143            callback=callback,1144            callback_steps=callback_steps,1145            **kwargs,1146        )1147