diffusers/community-pipelines-mirror

1import inspect2import re3from typing import Callable, List, Optional, Union4 5import numpy as np6import PIL.Image7import 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 437    if version.parse(version.parse(diffusers.__version__).base_version) >= version.parse("0.9.0"):438 439        def __init__(440            self,441            vae_encoder: OnnxRuntimeModel,442            vae_decoder: OnnxRuntimeModel,443            text_encoder: OnnxRuntimeModel,444            tokenizer: CLIPTokenizer,445            unet: OnnxRuntimeModel,446            scheduler: SchedulerMixin,447            safety_checker: OnnxRuntimeModel,448            feature_extractor: CLIPImageProcessor,449            requires_safety_checker: bool = True,450        ):451            super().__init__(452                vae_encoder=vae_encoder,453                vae_decoder=vae_decoder,454                text_encoder=text_encoder,455                tokenizer=tokenizer,456                unet=unet,457                scheduler=scheduler,458                safety_checker=safety_checker,459                feature_extractor=feature_extractor,460                requires_safety_checker=requires_safety_checker,461            )462            self.__init__additional__()463 464    else:465 466        def __init__(467            self,468            vae_encoder: OnnxRuntimeModel,469            vae_decoder: OnnxRuntimeModel,470            text_encoder: OnnxRuntimeModel,471            tokenizer: CLIPTokenizer,472            unet: OnnxRuntimeModel,473            scheduler: SchedulerMixin,474            safety_checker: OnnxRuntimeModel,475            feature_extractor: CLIPImageProcessor,476        ):477            super().__init__(478                vae_encoder=vae_encoder,479                vae_decoder=vae_decoder,480                text_encoder=text_encoder,481                tokenizer=tokenizer,482                unet=unet,483                scheduler=scheduler,484                safety_checker=safety_checker,485                feature_extractor=feature_extractor,486            )487            self.__init__additional__()488 489    def __init__additional__(self):490        self.unet.config.in_channels = 4491        self.vae_scale_factor = 8492 493    def _encode_prompt(494        self,495        prompt,496        num_images_per_prompt,497        do_classifier_free_guidance,498        negative_prompt,499        max_embeddings_multiples,500    ):501        r"""502        Encodes the prompt into text encoder hidden states.503 504        Args:505            prompt (`str` or `list(int)`):506                prompt to be encoded507            num_images_per_prompt (`int`):508                number of images that should be generated per prompt509            do_classifier_free_guidance (`bool`):510                whether to use classifier free guidance or not511            negative_prompt (`str` or `List[str]`):512                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored513                if `guidance_scale` is less than `1`).514            max_embeddings_multiples (`int`, *optional*, defaults to `3`):515                The max multiple length of prompt embeddings compared to the max output length of text encoder.516        """517        batch_size = len(prompt) if isinstance(prompt, list) else 1518 519        if negative_prompt is None:520            negative_prompt = [""] * batch_size521        elif isinstance(negative_prompt, str):522            negative_prompt = [negative_prompt] * batch_size523        if batch_size != len(negative_prompt):524            raise ValueError(525                f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"526                f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"527                " the batch size of `prompt`."528            )529 530        text_embeddings, uncond_embeddings = get_weighted_text_embeddings(531            pipe=self,532            prompt=prompt,533            uncond_prompt=negative_prompt if do_classifier_free_guidance else None,534            max_embeddings_multiples=max_embeddings_multiples,535        )536 537        text_embeddings = text_embeddings.repeat(num_images_per_prompt, 0)538        if do_classifier_free_guidance:539            uncond_embeddings = uncond_embeddings.repeat(num_images_per_prompt, 0)540            text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])541 542        return text_embeddings543 544    def check_inputs(self, prompt, height, width, strength, callback_steps):545        if not isinstance(prompt, str) and not isinstance(prompt, list):546            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")547 548        if strength < 0 or strength > 1:549            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")550 551        if height % 8 != 0 or width % 8 != 0:552            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")553 554        if (callback_steps is None) or (555            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)556        ):557            raise ValueError(558                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"559                f" {type(callback_steps)}."560            )561 562    def get_timesteps(self, num_inference_steps, strength, is_text2img):563        if is_text2img:564            return self.scheduler.timesteps, num_inference_steps565        else:566            # get the original timestep using init_timestep567            offset = self.scheduler.config.get("steps_offset", 0)568            init_timestep = int(num_inference_steps * strength) + offset569            init_timestep = min(init_timestep, num_inference_steps)570 571            t_start = max(num_inference_steps - init_timestep + offset, 0)572            timesteps = self.scheduler.timesteps[t_start:]573            return timesteps, num_inference_steps - t_start574 575    def run_safety_checker(self, image):576        if self.safety_checker is not None:577            safety_checker_input = self.feature_extractor(578                self.numpy_to_pil(image), return_tensors="np"579            ).pixel_values.astype(image.dtype)580            # There will throw an error if use safety_checker directly and batchsize>1581            images, has_nsfw_concept = [], []582            for i in range(image.shape[0]):583                image_i, has_nsfw_concept_i = self.safety_checker(584                    clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1]585                )586                images.append(image_i)587                has_nsfw_concept.append(has_nsfw_concept_i[0])588            image = np.concatenate(images)589        else:590            has_nsfw_concept = None591        return image, has_nsfw_concept592 593    def decode_latents(self, latents):594        latents = 1 / 0.18215 * latents595        # image = self.vae_decoder(latent_sample=latents)[0]596        # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1597        image = np.concatenate(598            [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]599        )600        image = np.clip(image / 2 + 0.5, 0, 1)601        image = image.transpose((0, 2, 3, 1))602        return image603 604    def prepare_extra_step_kwargs(self, generator, eta):605        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature606        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.607        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502608        # and should be between [0, 1]609 610        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())611        extra_step_kwargs = {}612        if accepts_eta:613            extra_step_kwargs["eta"] = eta614 615        # check if the scheduler accepts generator616        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())617        if accepts_generator:618            extra_step_kwargs["generator"] = generator619        return extra_step_kwargs620 621    def prepare_latents(self, image, timestep, batch_size, height, width, dtype, generator, latents=None):622        if image is None:623            shape = (624                batch_size,625                self.unet.config.in_channels,626                height // self.vae_scale_factor,627                width // self.vae_scale_factor,628            )629 630            if latents is None:631                latents = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)632            else:633                if latents.shape != shape:634                    raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")635 636            # scale the initial noise by the standard deviation required by the scheduler637            latents = (torch.from_numpy(latents) * self.scheduler.init_noise_sigma).numpy()638            return latents, None, None639        else:640            init_latents = self.vae_encoder(sample=image)[0]641            init_latents = 0.18215 * init_latents642            init_latents = np.concatenate([init_latents] * batch_size, axis=0)643            init_latents_orig = init_latents644            shape = init_latents.shape645 646            # add noise to latents using the timesteps647            noise = torch.randn(shape, generator=generator, device="cpu").numpy().astype(dtype)648            latents = self.scheduler.add_noise(649                torch.from_numpy(init_latents), torch.from_numpy(noise), timestep650            ).numpy()651            return latents, init_latents_orig, noise652 653    @torch.no_grad()654    def __call__(655        self,656        prompt: Union[str, List[str]],657        negative_prompt: Optional[Union[str, List[str]]] = None,658        image: Union[np.ndarray, PIL.Image.Image] = None,659        mask_image: Union[np.ndarray, PIL.Image.Image] = None,660        height: int = 512,661        width: int = 512,662        num_inference_steps: int = 50,663        guidance_scale: float = 7.5,664        strength: float = 0.8,665        num_images_per_prompt: Optional[int] = 1,666        eta: float = 0.0,667        generator: Optional[torch.Generator] = None,668        latents: Optional[np.ndarray] = None,669        max_embeddings_multiples: Optional[int] = 3,670        output_type: Optional[str] = "pil",671        return_dict: bool = True,672        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,673        is_cancelled_callback: Optional[Callable[[], bool]] = None,674        callback_steps: int = 1,675        **kwargs,676    ):677        r"""678        Function invoked when calling the pipeline for generation.679 680        Args:681            prompt (`str` or `List[str]`):682                The prompt or prompts to guide the image generation.683            negative_prompt (`str` or `List[str]`, *optional*):684                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored685                if `guidance_scale` is less than `1`).686            image (`np.ndarray` or `PIL.Image.Image`):687                `Image`, or tensor representing an image batch, that will be used as the starting point for the688                process.689            mask_image (`np.ndarray` or `PIL.Image.Image`):690                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be691                replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a692                PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should693                contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.694            height (`int`, *optional*, defaults to 512):695                The height in pixels of the generated image.696            width (`int`, *optional*, defaults to 512):697                The width in pixels of the generated image.698            num_inference_steps (`int`, *optional*, defaults to 50):699                The number of denoising steps. More denoising steps usually lead to a higher quality image at the700                expense of slower inference.701            guidance_scale (`float`, *optional*, defaults to 7.5):702                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).703                `guidance_scale` is defined as `w` of equation 2. of [Imagen704                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >705                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,706                usually at the expense of lower image quality.707            strength (`float`, *optional*, defaults to 0.8):708                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.709                `image` will be used as a starting point, adding more noise to it the larger the `strength`. The710                number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added711                noise will be maximum and the denoising process will run for the full number of iterations specified in712                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.713            num_images_per_prompt (`int`, *optional*, defaults to 1):714                The number of images to generate per prompt.715            eta (`float`, *optional*, defaults to 0.0):716                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to717                [`schedulers.DDIMScheduler`], will be ignored for others.718            generator (`torch.Generator`, *optional*):719                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation720                deterministic.721            latents (`np.ndarray`, *optional*):722                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image723                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents724                tensor will ge generated by sampling using the supplied random `generator`.725            max_embeddings_multiples (`int`, *optional*, defaults to `3`):726                The max multiple length of prompt embeddings compared to the max output length of text encoder.727            output_type (`str`, *optional*, defaults to `"pil"`):728                The output format of the generate image. Choose between729                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.730            return_dict (`bool`, *optional*, defaults to `True`):731                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a732                plain tuple.733            callback (`Callable`, *optional*):734                A function that will be called every `callback_steps` steps during inference. The function will be735                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.736            is_cancelled_callback (`Callable`, *optional*):737                A function that will be called every `callback_steps` steps during inference. If the function returns738                `True`, the inference will be cancelled.739            callback_steps (`int`, *optional*, defaults to 1):740                The frequency at which the `callback` function will be called. If not specified, the callback will be741                called at every step.742 743        Returns:744            `None` if cancelled by `is_cancelled_callback`,745            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:746            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.747            When returning a tuple, the first element is a list with the generated images, and the second element is a748            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"749            (nsfw) content, according to the `safety_checker`.750        """751        # 0. Default height and width to unet752        height = height or self.unet.config.sample_size * self.vae_scale_factor753        width = width or self.unet.config.sample_size * self.vae_scale_factor754 755        # 1. Check inputs. Raise error if not correct756        self.check_inputs(prompt, height, width, strength, callback_steps)757 758        # 2. Define call parameters759        batch_size = 1 if isinstance(prompt, str) else len(prompt)760        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)761        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`762        # corresponds to doing no classifier free guidance.763        do_classifier_free_guidance = guidance_scale > 1.0764 765        # 3. Encode input prompt766        text_embeddings = self._encode_prompt(767            prompt,768            num_images_per_prompt,769            do_classifier_free_guidance,770            negative_prompt,771            max_embeddings_multiples,772        )773        dtype = text_embeddings.dtype774 775        # 4. Preprocess image and mask776        if isinstance(image, PIL.Image.Image):777            image = preprocess_image(image)778        if image is not None:779            image = image.astype(dtype)780        if isinstance(mask_image, PIL.Image.Image):781            mask_image = preprocess_mask(mask_image, self.vae_scale_factor)782        if mask_image is not None:783            mask = mask_image.astype(dtype)784            mask = np.concatenate([mask] * batch_size * num_images_per_prompt)785        else:786            mask = None787 788        # 5. set timesteps789        self.scheduler.set_timesteps(num_inference_steps)790        timestep_dtype = next(791            (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)"792        )793        timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype]794        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, image is None)795        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)796 797        # 6. Prepare latent variables798        latents, init_latents_orig, noise = self.prepare_latents(799            image,800            latent_timestep,801            batch_size * num_images_per_prompt,802            height,803            width,804            dtype,805            generator,806            latents,807        )808 809        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline810        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)811 812        # 8. Denoising loop813        for i, t in enumerate(self.progress_bar(timesteps)):814            # expand the latents if we are doing classifier free guidance815            latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents816            latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)817            latent_model_input = latent_model_input.numpy()818 819            # predict the noise residual820            noise_pred = self.unet(821                sample=latent_model_input,822                timestep=np.array([t], dtype=timestep_dtype),823                encoder_hidden_states=text_embeddings,824            )825            noise_pred = noise_pred[0]826 827            # perform guidance828            if do_classifier_free_guidance:829                noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)830                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)831 832            # compute the previous noisy sample x_t -> x_t-1833            scheduler_output = self.scheduler.step(834                torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs835            )836            latents = scheduler_output.prev_sample.numpy()837 838            if mask is not None:839                # masking840                init_latents_proper = self.scheduler.add_noise(841                    torch.from_numpy(init_latents_orig),842                    torch.from_numpy(noise),843                    t,844                ).numpy()845                latents = (init_latents_proper * mask) + (latents * (1 - mask))846 847            # call the callback, if provided848            if i % callback_steps == 0:849                if callback is not None:850                    step_idx = i // getattr(self.scheduler, "order", 1)851                    callback(step_idx, t, latents)852                if is_cancelled_callback is not None and is_cancelled_callback():853                    return None854 855        # 9. Post-processing856        image = self.decode_latents(latents)857 858        # 10. Run safety checker859        image, has_nsfw_concept = self.run_safety_checker(image)860 861        # 11. Convert to PIL862        if output_type == "pil":863            image = self.numpy_to_pil(image)864 865        if not return_dict:866            return image, has_nsfw_concept867 868        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)869 870    def text2img(871        self,872        prompt: Union[str, List[str]],873        negative_prompt: Optional[Union[str, List[str]]] = None,874        height: int = 512,875        width: int = 512,876        num_inference_steps: int = 50,877        guidance_scale: float = 7.5,878        num_images_per_prompt: Optional[int] = 1,879        eta: float = 0.0,880        generator: Optional[torch.Generator] = None,881        latents: Optional[np.ndarray] = None,882        max_embeddings_multiples: Optional[int] = 3,883        output_type: Optional[str] = "pil",884        return_dict: bool = True,885        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,886        callback_steps: int = 1,887        **kwargs,888    ):889        r"""890        Function for text-to-image generation.891        Args:892            prompt (`str` or `List[str]`):893                The prompt or prompts to guide the image generation.894            negative_prompt (`str` or `List[str]`, *optional*):895                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored896                if `guidance_scale` is less than `1`).897            height (`int`, *optional*, defaults to 512):898                The height in pixels of the generated image.899            width (`int`, *optional*, defaults to 512):900                The width in pixels of the generated image.901            num_inference_steps (`int`, *optional*, defaults to 50):902                The number of denoising steps. More denoising steps usually lead to a higher quality image at the903                expense of slower inference.904            guidance_scale (`float`, *optional*, defaults to 7.5):905                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).906                `guidance_scale` is defined as `w` of equation 2. of [Imagen907                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >908                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,909                usually at the expense of lower image quality.910            num_images_per_prompt (`int`, *optional*, defaults to 1):911                The number of images to generate per prompt.912            eta (`float`, *optional*, defaults to 0.0):913                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to914                [`schedulers.DDIMScheduler`], will be ignored for others.915            generator (`torch.Generator`, *optional*):916                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation917                deterministic.918            latents (`np.ndarray`, *optional*):919                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image920                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents921                tensor will ge generated by sampling using the supplied random `generator`.922            max_embeddings_multiples (`int`, *optional*, defaults to `3`):923                The max multiple length of prompt embeddings compared to the max output length of text encoder.924            output_type (`str`, *optional*, defaults to `"pil"`):925                The output format of the generate image. Choose between926                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.927            return_dict (`bool`, *optional*, defaults to `True`):928                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a929                plain tuple.930            callback (`Callable`, *optional*):931                A function that will be called every `callback_steps` steps during inference. The function will be932                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.933            callback_steps (`int`, *optional*, defaults to 1):934                The frequency at which the `callback` function will be called. If not specified, the callback will be935                called at every step.936        Returns:937            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:938            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.939            When returning a tuple, the first element is a list with the generated images, and the second element is a940            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"941            (nsfw) content, according to the `safety_checker`.942        """943        return self.__call__(944            prompt=prompt,945            negative_prompt=negative_prompt,946            height=height,947            width=width,948            num_inference_steps=num_inference_steps,949            guidance_scale=guidance_scale,950            num_images_per_prompt=num_images_per_prompt,951            eta=eta,952            generator=generator,953            latents=latents,954            max_embeddings_multiples=max_embeddings_multiples,955            output_type=output_type,956            return_dict=return_dict,957            callback=callback,958            callback_steps=callback_steps,959            **kwargs,960        )961 962    def img2img(963        self,964        image: Union[np.ndarray, PIL.Image.Image],965        prompt: Union[str, List[str]],966        negative_prompt: Optional[Union[str, List[str]]] = None,967        strength: float = 0.8,968        num_inference_steps: Optional[int] = 50,969        guidance_scale: Optional[float] = 7.5,970        num_images_per_prompt: Optional[int] = 1,971        eta: Optional[float] = 0.0,972        generator: Optional[torch.Generator] = None,973        max_embeddings_multiples: Optional[int] = 3,974        output_type: Optional[str] = "pil",975        return_dict: bool = True,976        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,977        callback_steps: int = 1,978        **kwargs,979    ):980        r"""981        Function for image-to-image generation.982        Args:983            image (`np.ndarray` or `PIL.Image.Image`):984                `Image`, or ndarray representing an image batch, that will be used as the starting point for the985                process.986            prompt (`str` or `List[str]`):987                The prompt or prompts to guide the image generation.988            negative_prompt (`str` or `List[str]`, *optional*):989                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored990                if `guidance_scale` is less than `1`).991            strength (`float`, *optional*, defaults to 0.8):992                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1.993                `image` will be used as a starting point, adding more noise to it the larger the `strength`. The994                number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added995                noise will be maximum and the denoising process will run for the full number of iterations specified in996                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.997            num_inference_steps (`int`, *optional*, defaults to 50):998                The number of denoising steps. More denoising steps usually lead to a higher quality image at the999                expense of slower inference. This parameter will be modulated by `strength`.1000            guidance_scale (`float`, *optional*, defaults to 7.5):1001                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1002                `guidance_scale` is defined as `w` of equation 2. of [Imagen1003                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1004                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1005                usually at the expense of lower image quality.1006            num_images_per_prompt (`int`, *optional*, defaults to 1):1007                The number of images to generate per prompt.1008            eta (`float`, *optional*, defaults to 0.0):1009                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to1010                [`schedulers.DDIMScheduler`], will be ignored for others.1011            generator (`torch.Generator`, *optional*):1012                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation1013                deterministic.1014            max_embeddings_multiples (`int`, *optional*, defaults to `3`):1015                The max multiple length of prompt embeddings compared to the max output length of text encoder.1016            output_type (`str`, *optional*, defaults to `"pil"`):1017                The output format of the generate image. Choose between1018                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.1019            return_dict (`bool`, *optional*, defaults to `True`):1020                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a1021                plain tuple.1022            callback (`Callable`, *optional*):1023                A function that will be called every `callback_steps` steps during inference. The function will be1024                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.1025            callback_steps (`int`, *optional*, defaults to 1):1026                The frequency at which the `callback` function will be called. If not specified, the callback will be1027                called at every step.1028        Returns:1029            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:1030            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.1031            When returning a tuple, the first element is a list with the generated images, and the second element is a1032            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"1033            (nsfw) content, according to the `safety_checker`.1034        """1035        return self.__call__(1036            prompt=prompt,1037            negative_prompt=negative_prompt,1038            image=image,1039            num_inference_steps=num_inference_steps,1040            guidance_scale=guidance_scale,1041            strength=strength,1042            num_images_per_prompt=num_images_per_prompt,1043            eta=eta,1044            generator=generator,1045            max_embeddings_multiples=max_embeddings_multiples,1046            output_type=output_type,1047            return_dict=return_dict,1048            callback=callback,1049            callback_steps=callback_steps,1050            **kwargs,1051        )1052 1053    def inpaint(1054        self,1055        image: Union[np.ndarray, PIL.Image.Image],1056        mask_image: Union[np.ndarray, PIL.Image.Image],1057        prompt: Union[str, List[str]],1058        negative_prompt: Optional[Union[str, List[str]]] = None,1059        strength: float = 0.8,1060        num_inference_steps: Optional[int] = 50,1061        guidance_scale: Optional[float] = 7.5,1062        num_images_per_prompt: Optional[int] = 1,1063        eta: Optional[float] = 0.0,1064        generator: Optional[torch.Generator] = None,1065        max_embeddings_multiples: Optional[int] = 3,1066        output_type: Optional[str] = "pil",1067        return_dict: bool = True,1068        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,1069        callback_steps: int = 1,1070        **kwargs,1071    ):1072        r"""1073        Function for inpaint.1074        Args:1075            image (`np.ndarray` or `PIL.Image.Image`):1076                `Image`, or tensor representing an image batch, that will be used as the starting point for the1077                process. This is the image whose masked region will be inpainted.1078            mask_image (`np.ndarray` or `PIL.Image.Image`):1079                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be1080                replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a1081                PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should1082                contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`.1083            prompt (`str` or `List[str]`):1084                The prompt or prompts to guide the image generation.1085            negative_prompt (`str` or `List[str]`, *optional*):1086                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored1087                if `guidance_scale` is less than `1`).1088            strength (`float`, *optional*, defaults to 0.8):1089                Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength`1090                is 1, the denoising process will be run on the masked area for the full number of iterations specified1091                in `num_inference_steps`. `image` will be used as a reference for the masked area, adding more1092                noise to that region the larger the `strength`. If `strength` is 0, no inpainting will occur.1093            num_inference_steps (`int`, *optional*, defaults to 50):1094                The reference number of denoising steps. More denoising steps usually lead to a higher quality image at1095                the expense of slower inference. This parameter will be modulated by `strength`, as explained above.1096            guidance_scale (`float`, *optional*, defaults to 7.5):1097                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1098                `guidance_scale` is defined as `w` of equation 2. of [Imagen1099                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1100                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1101                usually at the expense of lower image quality.1102            num_images_per_prompt (`int`, *optional*, defaults to 1):1103                The number of images to generate per prompt.1104            eta (`float`, *optional*, defaults to 0.0):1105                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to1106                [`schedulers.DDIMScheduler`], will be ignored for others.1107            generator (`torch.Generator`, *optional*):1108                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation1109                deterministic.1110            max_embeddings_multiples (`int`, *optional*, defaults to `3`):1111                The max multiple length of prompt embeddings compared to the max output length of text encoder.1112            output_type (`str`, *optional*, defaults to `"pil"`):1113                The output format of the generate image. Choose between1114                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.1115            return_dict (`bool`, *optional*, defaults to `True`):1116                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a1117                plain tuple.1118            callback (`Callable`, *optional*):1119                A function that will be called every `callback_steps` steps during inference. The function will be1120                called with the following arguments: `callback(step: int, timestep: int, latents: np.ndarray)`.1121            callback_steps (`int`, *optional*, defaults to 1):1122                The frequency at which the `callback` function will be called. If not specified, the callback will be1123                called at every step.1124        Returns:1125            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:1126            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.1127            When returning a tuple, the first element is a list with the generated images, and the second element is a1128            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"1129            (nsfw) content, according to the `safety_checker`.1130        """1131        return self.__call__(1132            prompt=prompt,1133            negative_prompt=negative_prompt,1134            image=image,1135            mask_image=mask_image,1136            num_inference_steps=num_inference_steps,1137            guidance_scale=guidance_scale,1138            strength=strength,1139            num_images_per_prompt=num_images_per_prompt,1140            eta=eta,1141            generator=generator,1142            max_embeddings_multiples=max_embeddings_multiples,1143            output_type=output_type,1144            return_dict=return_dict,1145            callback=callback,1146            callback_steps=callback_steps,1147            **kwargs,1148        )1149