diffusers/community-pipelines-mirror

1# Copyright 2024 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import inspect16from typing import Callable, List, Optional, Union17 18import numpy as np19import PIL.Image20import torch21from packaging import version22from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer23 24from diffusers import AutoencoderKL, DiffusionPipeline, UNet2DConditionModel25from diffusers.configuration_utils import FrozenDict, deprecate26from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin27from diffusers.pipelines.pipeline_utils import StableDiffusionMixin28from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput29from diffusers.pipelines.stable_diffusion.safety_checker import (30    StableDiffusionSafetyChecker,31)32from diffusers.schedulers import KarrasDiffusionSchedulers33from diffusers.utils import (34    logging,35)36from diffusers.utils.torch_utils import randn_tensor37 38 39logger = logging.get_logger(__name__)  # pylint: disable=invalid-name40 41 42def prepare_mask_and_masked_image(image, mask):43    """44    Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be45    converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the46    ``image`` and ``1`` for the ``mask``.47    The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be48    binarized (``mask > 0.5``) and cast to ``torch.float32`` too.49    Args:50        image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.51            It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``52            ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.53        mask (_type_): The mask to apply to the image, i.e. regions to inpaint.54            It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width``55            ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``.56    Raises:57        ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask58        should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions.59        TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not60            (ot the other way around).61    Returns:62        tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 463            dimensions: ``batch x channels x height x width``.64    """65    if isinstance(image, torch.Tensor):66        if not isinstance(mask, torch.Tensor):67            raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not")68 69        # Batch single image70        if image.ndim == 3:71            assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)"72            image = image.unsqueeze(0)73 74        # Batch and add channel dim for single mask75        if mask.ndim == 2:76            mask = mask.unsqueeze(0).unsqueeze(0)77 78        # Batch single mask or add channel dim79        if mask.ndim == 3:80            # Single batched mask, no channel dim or single mask not batched but channel dim81            if mask.shape[0] == 1:82                mask = mask.unsqueeze(0)83 84            # Batched masks no channel dim85            else:86                mask = mask.unsqueeze(1)87 88        assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions"89        assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions"90        assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size"91 92        # Check image is in [-1, 1]93        if image.min() < -1 or image.max() > 1:94            raise ValueError("Image should be in [-1, 1] range")95 96        # Check mask is in [0, 1]97        if mask.min() < 0 or mask.max() > 1:98            raise ValueError("Mask should be in [0, 1] range")99 100        # Binarize mask101        mask[mask < 0.5] = 0102        mask[mask >= 0.5] = 1103 104        # Image as float32105        image = image.to(dtype=torch.float32)106    elif isinstance(mask, torch.Tensor):107        raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not")108    else:109        # preprocess image110        if isinstance(image, (PIL.Image.Image, np.ndarray)):111            image = [image]112 113        if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):114            image = [np.array(i.convert("RGB"))[None, :] for i in image]115            image = np.concatenate(image, axis=0)116        elif isinstance(image, list) and isinstance(image[0], np.ndarray):117            image = np.concatenate([i[None, :] for i in image], axis=0)118 119        image = image.transpose(0, 3, 1, 2)120        image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0121 122        # preprocess mask123        if isinstance(mask, (PIL.Image.Image, np.ndarray)):124            mask = [mask]125 126        if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):127            mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0)128            mask = mask.astype(np.float32) / 255.0129        elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):130            mask = np.concatenate([m[None, None, :] for m in mask], axis=0)131 132        mask[mask < 0.5] = 0133        mask[mask >= 0.5] = 1134        mask = torch.from_numpy(mask)135 136    # masked_image = image * (mask >= 0.5)137    masked_image = image138 139    return mask, masked_image140 141 142class StableDiffusionRepaintPipeline(143    DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin, LoraLoaderMixin144):145    r"""146    Pipeline for text-guided image inpainting using Stable Diffusion. *This is an experimental feature*.147    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the148    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)149    In addition the pipeline inherits the following loading methods:150        - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`]151        - *LoRA*: [`loaders.LoraLoaderMixin.load_lora_weights`]152    as well as the following saving methods:153        - *LoRA*: [`loaders.LoraLoaderMixin.save_lora_weights`]154    Args:155        vae ([`AutoencoderKL`]):156            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.157        text_encoder ([`CLIPTextModel`]):158            Frozen text-encoder. Stable Diffusion uses the text portion of159            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically160            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.161        tokenizer (`CLIPTokenizer`):162            Tokenizer of class163            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).164        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.165        scheduler ([`SchedulerMixin`]):166            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of167            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].168        safety_checker ([`StableDiffusionSafetyChecker`]):169            Classification module that estimates whether generated images could be considered offensive or harmful.170            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.171        feature_extractor ([`CLIPImageProcessor`]):172            Model that extracts features from generated images to be used as inputs for the `safety_checker`.173    """174 175    _optional_components = ["safety_checker", "feature_extractor"]176 177    def __init__(178        self,179        vae: AutoencoderKL,180        text_encoder: CLIPTextModel,181        tokenizer: CLIPTokenizer,182        unet: UNet2DConditionModel,183        scheduler: KarrasDiffusionSchedulers,184        safety_checker: StableDiffusionSafetyChecker,185        feature_extractor: CLIPImageProcessor,186        requires_safety_checker: bool = True,187    ):188        super().__init__()189 190        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:191            deprecation_message = (192                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"193                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "194                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"195                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"196                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"197                " file"198            )199            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)200            new_config = dict(scheduler.config)201            new_config["steps_offset"] = 1202            scheduler._internal_dict = FrozenDict(new_config)203 204        if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:205            deprecation_message = (206                f"The configuration file of this scheduler: {scheduler} has not set the configuration"207                " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"208                " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"209                " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"210                " Hub, it would be very nice if you could open a Pull request for the"211                " `scheduler/scheduler_config.json` file"212            )213            deprecate(214                "skip_prk_steps not set",215                "1.0.0",216                deprecation_message,217                standard_warn=False,218            )219            new_config = dict(scheduler.config)220            new_config["skip_prk_steps"] = True221            scheduler._internal_dict = FrozenDict(new_config)222 223        if safety_checker is None and requires_safety_checker:224            logger.warning(225                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"226                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"227                " results in services or applications open to the public. Both the diffusers team and Hugging Face"228                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"229                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"230                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."231            )232 233        if safety_checker is not None and feature_extractor is None:234            raise ValueError(235                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"236                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."237            )238 239        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(240            version.parse(unet.config._diffusers_version).base_version241        ) < version.parse("0.9.0.dev0")242        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64243        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:244            deprecation_message = (245                "The configuration file of the unet has set the default `sample_size` to smaller than"246                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"247                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"248                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"249                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"250                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"251                " in the config might lead to incorrect results in future versions. If you have downloaded this"252                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"253                " the `unet/config.json` file"254            )255            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)256            new_config = dict(unet.config)257            new_config["sample_size"] = 64258            unet._internal_dict = FrozenDict(new_config)259        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4260        if unet.config.in_channels != 4:261            logger.warning(262                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"263                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"264                ". If you did not intend to modify"265                " this behavior, please check whether you have loaded the right checkpoint."266            )267 268        self.register_modules(269            vae=vae,270            text_encoder=text_encoder,271            tokenizer=tokenizer,272            unet=unet,273            scheduler=scheduler,274            safety_checker=safety_checker,275            feature_extractor=feature_extractor,276        )277        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)278        self.register_to_config(requires_safety_checker=requires_safety_checker)279 280    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt281    def _encode_prompt(282        self,283        prompt,284        device,285        num_images_per_prompt,286        do_classifier_free_guidance,287        negative_prompt=None,288        prompt_embeds: Optional[torch.Tensor] = None,289        negative_prompt_embeds: Optional[torch.Tensor] = None,290    ):291        r"""292        Encodes the prompt into text encoder hidden states.293        Args:294             prompt (`str` or `List[str]`, *optional*):295                prompt to be encoded296            device: (`torch.device`):297                torch device298            num_images_per_prompt (`int`):299                number of images that should be generated per prompt300            do_classifier_free_guidance (`bool`):301                whether to use classifier free guidance or not302            negative_prompt (`str` or `List[str]`, *optional*):303                The prompt or prompts not to guide the image generation. If not defined, one has to pass304                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is305                less than `1`).306            prompt_embeds (`torch.Tensor`, *optional*):307                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not308                provided, text embeddings will be generated from `prompt` input argument.309            negative_prompt_embeds (`torch.Tensor`, *optional*):310                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt311                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input312                argument.313        """314        if prompt is not None and isinstance(prompt, str):315            batch_size = 1316        elif prompt is not None and isinstance(prompt, list):317            batch_size = len(prompt)318        else:319            batch_size = prompt_embeds.shape[0]320 321        if prompt_embeds is None:322            # textual inversion: process multi-vector tokens if necessary323            if isinstance(self, TextualInversionLoaderMixin):324                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)325 326            text_inputs = self.tokenizer(327                prompt,328                padding="max_length",329                max_length=self.tokenizer.model_max_length,330                truncation=True,331                return_tensors="pt",332            )333            text_input_ids = text_inputs.input_ids334            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids335 336            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(337                text_input_ids, untruncated_ids338            ):339                removed_text = self.tokenizer.batch_decode(340                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]341                )342                logger.warning(343                    "The following part of your input was truncated because CLIP can only handle sequences up to"344                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"345                )346 347            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:348                attention_mask = text_inputs.attention_mask.to(device)349            else:350                attention_mask = None351 352            prompt_embeds = self.text_encoder(353                text_input_ids.to(device),354                attention_mask=attention_mask,355            )356            prompt_embeds = prompt_embeds[0]357 358        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)359 360        bs_embed, seq_len, _ = prompt_embeds.shape361        # duplicate text embeddings for each generation per prompt, using mps friendly method362        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)363        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)364 365        # get unconditional embeddings for classifier free guidance366        if do_classifier_free_guidance and negative_prompt_embeds is None:367            uncond_tokens: List[str]368            if negative_prompt is None:369                uncond_tokens = [""] * batch_size370            elif type(prompt) is not type(negative_prompt):371                raise TypeError(372                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="373                    f" {type(prompt)}."374                )375            elif isinstance(negative_prompt, str):376                uncond_tokens = [negative_prompt]377            elif batch_size != len(negative_prompt):378                raise ValueError(379                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"380                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"381                    " the batch size of `prompt`."382                )383            else:384                uncond_tokens = negative_prompt385 386            # textual inversion: process multi-vector tokens if necessary387            if isinstance(self, TextualInversionLoaderMixin):388                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)389 390            max_length = prompt_embeds.shape[1]391            uncond_input = self.tokenizer(392                uncond_tokens,393                padding="max_length",394                max_length=max_length,395                truncation=True,396                return_tensors="pt",397            )398 399            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:400                attention_mask = uncond_input.attention_mask.to(device)401            else:402                attention_mask = None403 404            negative_prompt_embeds = self.text_encoder(405                uncond_input.input_ids.to(device),406                attention_mask=attention_mask,407            )408            negative_prompt_embeds = negative_prompt_embeds[0]409 410        if do_classifier_free_guidance:411            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method412            seq_len = negative_prompt_embeds.shape[1]413 414            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)415 416            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)417            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)418 419            # For classifier free guidance, we need to do two forward passes.420            # Here we concatenate the unconditional and text embeddings into a single batch421            # to avoid doing two forward passes422            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])423 424        return prompt_embeds425 426    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker427    def run_safety_checker(self, image, device, dtype):428        if self.safety_checker is not None:429            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)430            image, has_nsfw_concept = self.safety_checker(431                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)432            )433        else:434            has_nsfw_concept = None435        return image, has_nsfw_concept436 437    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs438    def prepare_extra_step_kwargs(self, generator, eta):439        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature440        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.441        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502442        # and should be between [0, 1]443 444        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())445        extra_step_kwargs = {}446        if accepts_eta:447            extra_step_kwargs["eta"] = eta448 449        # check if the scheduler accepts generator450        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())451        if accepts_generator:452            extra_step_kwargs["generator"] = generator453        return extra_step_kwargs454 455    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents456    def decode_latents(self, latents):457        latents = 1 / self.vae.config.scaling_factor * latents458        image = self.vae.decode(latents).sample459        image = (image / 2 + 0.5).clamp(0, 1)460        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16461        image = image.cpu().permute(0, 2, 3, 1).float().numpy()462        return image463 464    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs465    def check_inputs(466        self,467        prompt,468        height,469        width,470        callback_steps,471        negative_prompt=None,472        prompt_embeds=None,473        negative_prompt_embeds=None,474    ):475        if height % 8 != 0 or width % 8 != 0:476            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")477 478        if (callback_steps is None) or (479            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)480        ):481            raise ValueError(482                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"483                f" {type(callback_steps)}."484            )485 486        if prompt is not None and prompt_embeds is not None:487            raise ValueError(488                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"489                " only forward one of the two."490            )491        elif prompt is None and prompt_embeds is None:492            raise ValueError(493                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."494            )495        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):496            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")497 498        if negative_prompt is not None and negative_prompt_embeds is not None:499            raise ValueError(500                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"501                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."502            )503 504        if prompt_embeds is not None and negative_prompt_embeds is not None:505            if prompt_embeds.shape != negative_prompt_embeds.shape:506                raise ValueError(507                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"508                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"509                    f" {negative_prompt_embeds.shape}."510                )511 512    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents513    def prepare_latents(514        self,515        batch_size,516        num_channels_latents,517        height,518        width,519        dtype,520        device,521        generator,522        latents=None,523    ):524        shape = (525            batch_size,526            num_channels_latents,527            height // self.vae_scale_factor,528            width // self.vae_scale_factor,529        )530        if isinstance(generator, list) and len(generator) != batch_size:531            raise ValueError(532                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"533                f" size of {batch_size}. Make sure the batch size matches the length of the generators."534            )535 536        if latents is None:537            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)538        else:539            latents = latents.to(device)540 541        # scale the initial noise by the standard deviation required by the scheduler542        latents = latents * self.scheduler.init_noise_sigma543        return latents544 545    def prepare_mask_latents(546        self,547        mask,548        masked_image,549        batch_size,550        height,551        width,552        dtype,553        device,554        generator,555        do_classifier_free_guidance,556    ):557        # resize the mask to latents shape as we concatenate the mask to the latents558        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload559        # and half precision560        mask = torch.nn.functional.interpolate(561            mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)562        )563        mask = mask.to(device=device, dtype=dtype)564 565        masked_image = masked_image.to(device=device, dtype=dtype)566 567        # encode the mask image into latents space so we can concatenate it to the latents568        if isinstance(generator, list):569            masked_image_latents = [570                self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])571                for i in range(batch_size)572            ]573            masked_image_latents = torch.cat(masked_image_latents, dim=0)574        else:575            masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)576        masked_image_latents = self.vae.config.scaling_factor * masked_image_latents577 578        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method579        if mask.shape[0] < batch_size:580            if not batch_size % mask.shape[0] == 0:581                raise ValueError(582                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"583                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"584                    " of masks that you pass is divisible by the total requested batch size."585                )586            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)587        if masked_image_latents.shape[0] < batch_size:588            if not batch_size % masked_image_latents.shape[0] == 0:589                raise ValueError(590                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"591                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."592                    " Make sure the number of images that you pass is divisible by the total requested batch size."593                )594            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)595 596        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask597        masked_image_latents = (598            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents599        )600 601        # aligning device to prevent device errors when concating it with the latent model input602        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)603        return mask, masked_image_latents604 605    @torch.no_grad()606    def __call__(607        self,608        prompt: Union[str, List[str]] = None,609        image: Union[torch.Tensor, PIL.Image.Image] = None,610        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,611        height: Optional[int] = None,612        width: Optional[int] = None,613        num_inference_steps: int = 50,614        jump_length: Optional[int] = 10,615        jump_n_sample: Optional[int] = 10,616        guidance_scale: float = 7.5,617        negative_prompt: Optional[Union[str, List[str]]] = None,618        num_images_per_prompt: Optional[int] = 1,619        eta: float = 0.0,620        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,621        latents: Optional[torch.Tensor] = None,622        prompt_embeds: Optional[torch.Tensor] = None,623        negative_prompt_embeds: Optional[torch.Tensor] = None,624        output_type: Optional[str] = "pil",625        return_dict: bool = True,626        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,627        callback_steps: int = 1,628    ):629        r"""630        Function invoked when calling the pipeline for generation.631        Args:632            prompt (`str` or `List[str]`, *optional*):633                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.634                instead.635            image (`PIL.Image.Image`):636                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will637                be masked out with `mask_image` and repainted according to `prompt`.638            mask_image (`PIL.Image.Image`):639                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be640                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted641                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)642                instead of 3, so the expected shape would be `(B, H, W, 1)`.643            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):644                The height in pixels of the generated image.645            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):646                The width in pixels of the generated image.647            num_inference_steps (`int`, *optional*, defaults to 50):648                The number of denoising steps. More denoising steps usually lead to a higher quality image at the649                expense of slower inference.650            jump_length (`int`, *optional*, defaults to 10):651                The number of steps taken forward in time before going backward in time for a single jump ("j" in652                RePaint paper). Take a look at Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.653            jump_n_sample (`int`, *optional*, defaults to 10):654                The number of times we will make forward time jump for a given chosen time sample. Take a look at655                Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.656            guidance_scale (`float`, *optional*, defaults to 7.5):657                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).658                `guidance_scale` is defined as `w` of equation 2. of [Imagen659                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >660                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,661                usually at the expense of lower image quality.662            negative_prompt (`str` or `List[str]`, *optional*):663                The prompt or prompts not to guide the image generation. If not defined, one has to pass664                `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`665                is less than `1`).666            num_images_per_prompt (`int`, *optional*, defaults to 1):667                The number of images to generate per prompt.668            eta (`float`, *optional*, defaults to 0.0):669                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to670                [`schedulers.DDIMScheduler`], will be ignored for others.671            generator (`torch.Generator`, *optional*):672                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)673                to make generation deterministic.674            latents (`torch.Tensor`, *optional*):675                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image676                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents677                tensor will ge generated by sampling using the supplied random `generator`.678            prompt_embeds (`torch.Tensor`, *optional*):679                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not680                provided, text embeddings will be generated from `prompt` input argument.681            negative_prompt_embeds (`torch.Tensor`, *optional*):682                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt683                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input684                argument.685            output_type (`str`, *optional*, defaults to `"pil"`):686                The output format of the generate image. Choose between687                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.688            return_dict (`bool`, *optional*, defaults to `True`):689                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a690                plain tuple.691            callback (`Callable`, *optional*):692                A function that will be called every `callback_steps` steps during inference. The function will be693                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.694            callback_steps (`int`, *optional*, defaults to 1):695                The frequency at which the `callback` function will be called. If not specified, the callback will be696                called at every step.697        Examples:698        ```py699        >>> import PIL700        >>> import requests701        >>> import torch702        >>> from io import BytesIO703        >>> from diffusers import StableDiffusionPipeline, RePaintScheduler704        >>> def download_image(url):705        ...     response = requests.get(url)706        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")707        >>> base_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/"708        >>> img_url = base_url + "overture-creations-5sI6fQgYIuo.png"709        >>> mask_url = base_url + "overture-creations-5sI6fQgYIuo_mask.png "710        >>> init_image = download_image(img_url).resize((512, 512))711        >>> mask_image = download_image(mask_url).resize((512, 512))712        >>> pipe = DiffusionPipeline.from_pretrained(713        ...     "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16, custom_pipeline="stable_diffusion_repaint",714        ... )715        >>> pipe.scheduler = RePaintScheduler.from_config(pipe.scheduler.config)716        >>> pipe = pipe.to("cuda")717        >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"718        >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]719        ```720        Returns:721            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:722            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.723            When returning a tuple, the first element is a list with the generated images, and the second element is a724            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"725            (nsfw) content, according to the `safety_checker`.726        """727        # 0. Default height and width to unet728        height = height or self.unet.config.sample_size * self.vae_scale_factor729        width = width or self.unet.config.sample_size * self.vae_scale_factor730 731        # 1. Check inputs732        self.check_inputs(733            prompt,734            height,735            width,736            callback_steps,737            negative_prompt,738            prompt_embeds,739            negative_prompt_embeds,740        )741 742        if image is None:743            raise ValueError("`image` input cannot be undefined.")744 745        if mask_image is None:746            raise ValueError("`mask_image` input cannot be undefined.")747 748        # 2. Define call parameters749        if prompt is not None and isinstance(prompt, str):750            batch_size = 1751        elif prompt is not None and isinstance(prompt, list):752            batch_size = len(prompt)753        else:754            batch_size = prompt_embeds.shape[0]755 756        device = self._execution_device757        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)758        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`759        # corresponds to doing no classifier free guidance.760        do_classifier_free_guidance = guidance_scale > 1.0761 762        # 3. Encode input prompt763        prompt_embeds = self._encode_prompt(764            prompt,765            device,766            num_images_per_prompt,767            do_classifier_free_guidance,768            negative_prompt,769            prompt_embeds=prompt_embeds,770            negative_prompt_embeds=negative_prompt_embeds,771        )772 773        # 4. Preprocess mask and image774        mask, masked_image = prepare_mask_and_masked_image(image, mask_image)775 776        # 5. set timesteps777        self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, device)778        self.scheduler.eta = eta779 780        timesteps = self.scheduler.timesteps781        # latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)782 783        # 6. Prepare latent variables784        num_channels_latents = self.vae.config.latent_channels785        latents = self.prepare_latents(786            batch_size * num_images_per_prompt,787            num_channels_latents,788            height,789            width,790            prompt_embeds.dtype,791            device,792            generator,793            latents,794        )795 796        # 7. Prepare mask latent variables797        mask, masked_image_latents = self.prepare_mask_latents(798            mask,799            masked_image,800            batch_size * num_images_per_prompt,801            height,802            width,803            prompt_embeds.dtype,804            device,805            generator,806            do_classifier_free_guidance=False,  # We do not need duplicate mask and image807        )808 809        # 8. Check that sizes of mask, masked image and latents match810        # num_channels_mask = mask.shape[1]811        # num_channels_masked_image = masked_image_latents.shape[1]812        if num_channels_latents != self.unet.config.in_channels:813            raise ValueError(814                f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"815                f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} "816                f" = Please verify the config of"817                " `pipeline.unet` or your `mask_image` or `image` input."818            )819 820        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline821        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)822 823        t_last = timesteps[0] + 1824 825        # 10. Denoising loop826        with self.progress_bar(total=len(timesteps)) as progress_bar:827            for i, t in enumerate(timesteps):828                if t >= t_last:829                    # compute the reverse: x_t-1 -> x_t830                    latents = self.scheduler.undo_step(latents, t_last, generator)831                    progress_bar.update()832                    t_last = t833                    continue834 835                # expand the latents if we are doing classifier free guidance836                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents837 838                # concat latents, mask, masked_image_latents in the channel dimension839                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)840                # latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)841 842                # predict the noise residual843                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample844 845                # perform guidance846                if do_classifier_free_guidance:847                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)848                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)849 850                # compute the previous noisy sample x_t -> x_t-1851                latents = self.scheduler.step(852                    noise_pred,853                    t,854                    latents,855                    masked_image_latents,856                    mask,857                    **extra_step_kwargs,858                ).prev_sample859 860                # call the callback, if provided861                progress_bar.update()862                if callback is not None and i % callback_steps == 0:863                    step_idx = i // getattr(self.scheduler, "order", 1)864                    callback(step_idx, t, latents)865 866                t_last = t867 868        # 11. Post-processing869        image = self.decode_latents(latents)870 871        # 12. Run safety checker872        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)873 874        # 13. Convert to PIL875        if output_type == "pil":876            image = self.numpy_to_pil(image)877 878        # Offload last model to CPU879        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:880            self.final_offload_hook.offload()881 882        if not return_dict:883            return (image, has_nsfw_concept)884 885        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)886