diffusers/community-pipelines-mirror

1# Copyright 2025 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 StableDiffusionLoraLoaderMixin, 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, StableDiffusionLoraLoaderMixin144):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.StableDiffusionLoraLoaderMixin.load_lora_weights`]152    as well as the following saving methods:153        - *LoRA*: [`loaders.StableDiffusionLoraLoaderMixin.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 scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 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 scheduler is not None and getattr(scheduler.config, "skip_prk_steps", True) 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 = (240            unet is not None241            and hasattr(unet.config, "_diffusers_version")242            and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0")243        )244        is_unet_sample_size_less_64 = (245            unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64246        )247        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:248            deprecation_message = (249                "The configuration file of the unet has set the default `sample_size` to smaller than"250                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"251                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"252                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"253                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"254                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"255                " in the config might lead to incorrect results in future versions. If you have downloaded this"256                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"257                " the `unet/config.json` file"258            )259            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)260            new_config = dict(unet.config)261            new_config["sample_size"] = 64262            unet._internal_dict = FrozenDict(new_config)263        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4264        if unet is not None and unet.config.in_channels != 4:265            logger.warning(266                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"267                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"268                ". If you did not intend to modify"269                " this behavior, please check whether you have loaded the right checkpoint."270            )271 272        self.register_modules(273            vae=vae,274            text_encoder=text_encoder,275            tokenizer=tokenizer,276            unet=unet,277            scheduler=scheduler,278            safety_checker=safety_checker,279            feature_extractor=feature_extractor,280        )281        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8282        self.register_to_config(requires_safety_checker=requires_safety_checker)283 284    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt285    def _encode_prompt(286        self,287        prompt,288        device,289        num_images_per_prompt,290        do_classifier_free_guidance,291        negative_prompt=None,292        prompt_embeds: Optional[torch.Tensor] = None,293        negative_prompt_embeds: Optional[torch.Tensor] = None,294    ):295        r"""296        Encodes the prompt into text encoder hidden states.297        Args:298             prompt (`str` or `List[str]`, *optional*):299                prompt to be encoded300            device: (`torch.device`):301                torch device302            num_images_per_prompt (`int`):303                number of images that should be generated per prompt304            do_classifier_free_guidance (`bool`):305                whether to use classifier free guidance or not306            negative_prompt (`str` or `List[str]`, *optional*):307                The prompt or prompts not to guide the image generation. If not defined, one has to pass308                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is309                less than `1`).310            prompt_embeds (`torch.Tensor`, *optional*):311                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not312                provided, text embeddings will be generated from `prompt` input argument.313            negative_prompt_embeds (`torch.Tensor`, *optional*):314                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt315                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input316                argument.317        """318        if prompt is not None and isinstance(prompt, str):319            batch_size = 1320        elif prompt is not None and isinstance(prompt, list):321            batch_size = len(prompt)322        else:323            batch_size = prompt_embeds.shape[0]324 325        if prompt_embeds is None:326            # textual inversion: process multi-vector tokens if necessary327            if isinstance(self, TextualInversionLoaderMixin):328                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)329 330            text_inputs = self.tokenizer(331                prompt,332                padding="max_length",333                max_length=self.tokenizer.model_max_length,334                truncation=True,335                return_tensors="pt",336            )337            text_input_ids = text_inputs.input_ids338            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids339 340            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(341                text_input_ids, untruncated_ids342            ):343                removed_text = self.tokenizer.batch_decode(344                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]345                )346                logger.warning(347                    "The following part of your input was truncated because CLIP can only handle sequences up to"348                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"349                )350 351            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:352                attention_mask = text_inputs.attention_mask.to(device)353            else:354                attention_mask = None355 356            prompt_embeds = self.text_encoder(357                text_input_ids.to(device),358                attention_mask=attention_mask,359            )360            prompt_embeds = prompt_embeds[0]361 362        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)363 364        bs_embed, seq_len, _ = prompt_embeds.shape365        # duplicate text embeddings for each generation per prompt, using mps friendly method366        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)367        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)368 369        # get unconditional embeddings for classifier free guidance370        if do_classifier_free_guidance and negative_prompt_embeds is None:371            uncond_tokens: List[str]372            if negative_prompt is None:373                uncond_tokens = [""] * batch_size374            elif type(prompt) is not type(negative_prompt):375                raise TypeError(376                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="377                    f" {type(prompt)}."378                )379            elif isinstance(negative_prompt, str):380                uncond_tokens = [negative_prompt]381            elif batch_size != len(negative_prompt):382                raise ValueError(383                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"384                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"385                    " the batch size of `prompt`."386                )387            else:388                uncond_tokens = negative_prompt389 390            # textual inversion: process multi-vector tokens if necessary391            if isinstance(self, TextualInversionLoaderMixin):392                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)393 394            max_length = prompt_embeds.shape[1]395            uncond_input = self.tokenizer(396                uncond_tokens,397                padding="max_length",398                max_length=max_length,399                truncation=True,400                return_tensors="pt",401            )402 403            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:404                attention_mask = uncond_input.attention_mask.to(device)405            else:406                attention_mask = None407 408            negative_prompt_embeds = self.text_encoder(409                uncond_input.input_ids.to(device),410                attention_mask=attention_mask,411            )412            negative_prompt_embeds = negative_prompt_embeds[0]413 414        if do_classifier_free_guidance:415            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method416            seq_len = negative_prompt_embeds.shape[1]417 418            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)419 420            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)421            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)422 423            # For classifier free guidance, we need to do two forward passes.424            # Here we concatenate the unconditional and text embeddings into a single batch425            # to avoid doing two forward passes426            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])427 428        return prompt_embeds429 430    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker431    def run_safety_checker(self, image, device, dtype):432        if self.safety_checker is not None:433            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)434            image, has_nsfw_concept = self.safety_checker(435                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)436            )437        else:438            has_nsfw_concept = None439        return image, has_nsfw_concept440 441    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs442    def prepare_extra_step_kwargs(self, generator, eta):443        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature444        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.445        # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502446        # and should be between [0, 1]447 448        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())449        extra_step_kwargs = {}450        if accepts_eta:451            extra_step_kwargs["eta"] = eta452 453        # check if the scheduler accepts generator454        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())455        if accepts_generator:456            extra_step_kwargs["generator"] = generator457        return extra_step_kwargs458 459    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents460    def decode_latents(self, latents):461        latents = 1 / self.vae.config.scaling_factor * latents462        image = self.vae.decode(latents).sample463        image = (image / 2 + 0.5).clamp(0, 1)464        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16465        image = image.cpu().permute(0, 2, 3, 1).float().numpy()466        return image467 468    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs469    def check_inputs(470        self,471        prompt,472        height,473        width,474        callback_steps,475        negative_prompt=None,476        prompt_embeds=None,477        negative_prompt_embeds=None,478    ):479        if height % 8 != 0 or width % 8 != 0:480            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")481 482        if (callback_steps is None) or (483            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)484        ):485            raise ValueError(486                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"487                f" {type(callback_steps)}."488            )489 490        if prompt is not None and prompt_embeds is not None:491            raise ValueError(492                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"493                " only forward one of the two."494            )495        elif prompt is None and prompt_embeds is None:496            raise ValueError(497                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."498            )499        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):500            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")501 502        if negative_prompt is not None and negative_prompt_embeds is not None:503            raise ValueError(504                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"505                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."506            )507 508        if prompt_embeds is not None and negative_prompt_embeds is not None:509            if prompt_embeds.shape != negative_prompt_embeds.shape:510                raise ValueError(511                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"512                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"513                    f" {negative_prompt_embeds.shape}."514                )515 516    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents517    def prepare_latents(518        self,519        batch_size,520        num_channels_latents,521        height,522        width,523        dtype,524        device,525        generator,526        latents=None,527    ):528        shape = (529            batch_size,530            num_channels_latents,531            height // self.vae_scale_factor,532            width // self.vae_scale_factor,533        )534        if isinstance(generator, list) and len(generator) != batch_size:535            raise ValueError(536                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"537                f" size of {batch_size}. Make sure the batch size matches the length of the generators."538            )539 540        if latents is None:541            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)542        else:543            latents = latents.to(device)544 545        # scale the initial noise by the standard deviation required by the scheduler546        latents = latents * self.scheduler.init_noise_sigma547        return latents548 549    def prepare_mask_latents(550        self,551        mask,552        masked_image,553        batch_size,554        height,555        width,556        dtype,557        device,558        generator,559        do_classifier_free_guidance,560    ):561        # resize the mask to latents shape as we concatenate the mask to the latents562        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload563        # and half precision564        mask = torch.nn.functional.interpolate(565            mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)566        )567        mask = mask.to(device=device, dtype=dtype)568 569        masked_image = masked_image.to(device=device, dtype=dtype)570 571        # encode the mask image into latents space so we can concatenate it to the latents572        if isinstance(generator, list):573            masked_image_latents = [574                self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])575                for i in range(batch_size)576            ]577            masked_image_latents = torch.cat(masked_image_latents, dim=0)578        else:579            masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)580        masked_image_latents = self.vae.config.scaling_factor * masked_image_latents581 582        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method583        if mask.shape[0] < batch_size:584            if not batch_size % mask.shape[0] == 0:585                raise ValueError(586                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"587                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"588                    " of masks that you pass is divisible by the total requested batch size."589                )590            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)591        if masked_image_latents.shape[0] < batch_size:592            if not batch_size % masked_image_latents.shape[0] == 0:593                raise ValueError(594                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"595                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."596                    " Make sure the number of images that you pass is divisible by the total requested batch size."597                )598            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)599 600        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask601        masked_image_latents = (602            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents603        )604 605        # aligning device to prevent device errors when concating it with the latent model input606        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)607        return mask, masked_image_latents608 609    @torch.no_grad()610    def __call__(611        self,612        prompt: Union[str, List[str]] = None,613        image: Union[torch.Tensor, PIL.Image.Image] = None,614        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,615        height: Optional[int] = None,616        width: Optional[int] = None,617        num_inference_steps: int = 50,618        jump_length: Optional[int] = 10,619        jump_n_sample: Optional[int] = 10,620        guidance_scale: float = 7.5,621        negative_prompt: Optional[Union[str, List[str]]] = None,622        num_images_per_prompt: Optional[int] = 1,623        eta: float = 0.0,624        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,625        latents: Optional[torch.Tensor] = None,626        prompt_embeds: Optional[torch.Tensor] = None,627        negative_prompt_embeds: Optional[torch.Tensor] = None,628        output_type: Optional[str] = "pil",629        return_dict: bool = True,630        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,631        callback_steps: int = 1,632    ):633        r"""634        Function invoked when calling the pipeline for generation.635        Args:636            prompt (`str` or `List[str]`, *optional*):637                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.638                instead.639            image (`PIL.Image.Image`):640                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will641                be masked out with `mask_image` and repainted according to `prompt`.642            mask_image (`PIL.Image.Image`):643                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be644                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted645                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)646                instead of 3, so the expected shape would be `(B, H, W, 1)`.647            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):648                The height in pixels of the generated image.649            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):650                The width in pixels of the generated image.651            num_inference_steps (`int`, *optional*, defaults to 50):652                The number of denoising steps. More denoising steps usually lead to a higher quality image at the653                expense of slower inference.654            jump_length (`int`, *optional*, defaults to 10):655                The number of steps taken forward in time before going backward in time for a single jump ("j" in656                RePaint paper). Take a look at Figure 9 and 10 in https://huggingface.co/papers/2201.09865.657            jump_n_sample (`int`, *optional*, defaults to 10):658                The number of times we will make forward time jump for a given chosen time sample. Take a look at659                Figure 9 and 10 in https://huggingface.co/papers/2201.09865.660            guidance_scale (`float`, *optional*, defaults to 7.5):661                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).662                `guidance_scale` is defined as `w` of equation 2. of [Imagen663                Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >664                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,665                usually at the expense of lower image quality.666            negative_prompt (`str` or `List[str]`, *optional*):667                The prompt or prompts not to guide the image generation. If not defined, one has to pass668                `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`669                is less than `1`).670            num_images_per_prompt (`int`, *optional*, defaults to 1):671                The number of images to generate per prompt.672            eta (`float`, *optional*, defaults to 0.0):673                Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to674                [`schedulers.DDIMScheduler`], will be ignored for others.675            generator (`torch.Generator`, *optional*):676                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)677                to make generation deterministic.678            latents (`torch.Tensor`, *optional*):679                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image680                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents681                tensor will ge generated by sampling using the supplied random `generator`.682            prompt_embeds (`torch.Tensor`, *optional*):683                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not684                provided, text embeddings will be generated from `prompt` input argument.685            negative_prompt_embeds (`torch.Tensor`, *optional*):686                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt687                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input688                argument.689            output_type (`str`, *optional*, defaults to `"pil"`):690                The output format of the generate image. Choose between691                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.692            return_dict (`bool`, *optional*, defaults to `True`):693                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a694                plain tuple.695            callback (`Callable`, *optional*):696                A function that will be called every `callback_steps` steps during inference. The function will be697                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.698            callback_steps (`int`, *optional*, defaults to 1):699                The frequency at which the `callback` function will be called. If not specified, the callback will be700                called at every step.701        Examples:702        ```py703        >>> import PIL704        >>> import requests705        >>> import torch706        >>> from io import BytesIO707        >>> from diffusers import StableDiffusionPipeline, RePaintScheduler708        >>> def download_image(url):709        ...     response = requests.get(url)710        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")711        >>> base_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/"712        >>> img_url = base_url + "overture-creations-5sI6fQgYIuo.png"713        >>> mask_url = base_url + "overture-creations-5sI6fQgYIuo_mask.png "714        >>> init_image = download_image(img_url).resize((512, 512))715        >>> mask_image = download_image(mask_url).resize((512, 512))716        >>> pipe = DiffusionPipeline.from_pretrained(717        ...     "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16, custom_pipeline="stable_diffusion_repaint",718        ... )719        >>> pipe.scheduler = RePaintScheduler.from_config(pipe.scheduler.config)720        >>> pipe = pipe.to("cuda")721        >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"722        >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]723        ```724        Returns:725            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:726            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.727            When returning a tuple, the first element is a list with the generated images, and the second element is a728            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"729            (nsfw) content, according to the `safety_checker`.730        """731        # 0. Default height and width to unet732        height = height or self.unet.config.sample_size * self.vae_scale_factor733        width = width or self.unet.config.sample_size * self.vae_scale_factor734 735        # 1. Check inputs736        self.check_inputs(737            prompt,738            height,739            width,740            callback_steps,741            negative_prompt,742            prompt_embeds,743            negative_prompt_embeds,744        )745 746        if image is None:747            raise ValueError("`image` input cannot be undefined.")748 749        if mask_image is None:750            raise ValueError("`mask_image` input cannot be undefined.")751 752        # 2. Define call parameters753        if prompt is not None and isinstance(prompt, str):754            batch_size = 1755        elif prompt is not None and isinstance(prompt, list):756            batch_size = len(prompt)757        else:758            batch_size = prompt_embeds.shape[0]759 760        device = self._execution_device761        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)762        # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`763        # corresponds to doing no classifier free guidance.764        do_classifier_free_guidance = guidance_scale > 1.0765 766        # 3. Encode input prompt767        prompt_embeds = self._encode_prompt(768            prompt,769            device,770            num_images_per_prompt,771            do_classifier_free_guidance,772            negative_prompt,773            prompt_embeds=prompt_embeds,774            negative_prompt_embeds=negative_prompt_embeds,775        )776 777        # 4. Preprocess mask and image778        mask, masked_image = prepare_mask_and_masked_image(image, mask_image)779 780        # 5. set timesteps781        self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, device)782        self.scheduler.eta = eta783 784        timesteps = self.scheduler.timesteps785        # latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)786 787        # 6. Prepare latent variables788        num_channels_latents = self.vae.config.latent_channels789        latents = self.prepare_latents(790            batch_size * num_images_per_prompt,791            num_channels_latents,792            height,793            width,794            prompt_embeds.dtype,795            device,796            generator,797            latents,798        )799 800        # 7. Prepare mask latent variables801        mask, masked_image_latents = self.prepare_mask_latents(802            mask,803            masked_image,804            batch_size * num_images_per_prompt,805            height,806            width,807            prompt_embeds.dtype,808            device,809            generator,810            do_classifier_free_guidance=False,  # We do not need duplicate mask and image811        )812 813        # 8. Check that sizes of mask, masked image and latents match814        # num_channels_mask = mask.shape[1]815        # num_channels_masked_image = masked_image_latents.shape[1]816        if num_channels_latents != self.unet.config.in_channels:817            raise ValueError(818                f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"819                f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} "820                f" = Please verify the config of"821                " `pipeline.unet` or your `mask_image` or `image` input."822            )823 824        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline825        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)826 827        t_last = timesteps[0] + 1828 829        # 10. Denoising loop830        with self.progress_bar(total=len(timesteps)) as progress_bar:831            for i, t in enumerate(timesteps):832                if t >= t_last:833                    # compute the reverse: x_t-1 -> x_t834                    latents = self.scheduler.undo_step(latents, t_last, generator)835                    progress_bar.update()836                    t_last = t837                    continue838 839                # expand the latents if we are doing classifier free guidance840                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents841 842                # concat latents, mask, masked_image_latents in the channel dimension843                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)844                # latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)845 846                # predict the noise residual847                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample848 849                # perform guidance850                if do_classifier_free_guidance:851                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)852                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)853 854                # compute the previous noisy sample x_t -> x_t-1855                latents = self.scheduler.step(856                    noise_pred,857                    t,858                    latents,859                    masked_image_latents,860                    mask,861                    **extra_step_kwargs,862                ).prev_sample863 864                # call the callback, if provided865                progress_bar.update()866                if callback is not None and i % callback_steps == 0:867                    step_idx = i // getattr(self.scheduler, "order", 1)868                    callback(step_idx, t, latents)869 870                t_last = t871 872        # 11. Post-processing873        image = self.decode_latents(latents)874 875        # 12. Run safety checker876        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)877 878        # 13. Convert to PIL879        if output_type == "pil":880            image = self.numpy_to_pil(image)881 882        # Offload last model to CPU883        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:884            self.final_offload_hook.offload()885 886        if not return_dict:887            return (image, has_nsfw_concept)888 889        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)890