diffusers/community-pipelines-mirror

1# Copyright 2023 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.stable_diffusion import StableDiffusionPipelineOutput28from diffusers.pipelines.stable_diffusion.safety_checker import (29    StableDiffusionSafetyChecker,30)31from diffusers.schedulers import KarrasDiffusionSchedulers32from diffusers.utils import (33    is_accelerate_available,34    is_accelerate_version,35    logging,36)37from diffusers.utils.torch_utils import randn_tensor38 39 40logger = logging.get_logger(__name__)  # pylint: disable=invalid-name41 42 43def prepare_mask_and_masked_image(image, mask):44    """45    Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be46    converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the47    ``image`` and ``1`` for the ``mask``.48    The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be49    binarized (``mask > 0.5``) and cast to ``torch.float32`` too.50    Args:51        image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.52            It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``53            ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.54        mask (_type_): The mask to apply to the image, i.e. regions to inpaint.55            It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width``56            ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``.57    Raises:58        ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask59        should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions.60        TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not61            (ot the other way around).62    Returns:63        tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 464            dimensions: ``batch x channels x height x width``.65    """66    if isinstance(image, torch.Tensor):67        if not isinstance(mask, torch.Tensor):68            raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not")69 70        # Batch single image71        if image.ndim == 3:72            assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)"73            image = image.unsqueeze(0)74 75        # Batch and add channel dim for single mask76        if mask.ndim == 2:77            mask = mask.unsqueeze(0).unsqueeze(0)78 79        # Batch single mask or add channel dim80        if mask.ndim == 3:81            # Single batched mask, no channel dim or single mask not batched but channel dim82            if mask.shape[0] == 1:83                mask = mask.unsqueeze(0)84 85            # Batched masks no channel dim86            else:87                mask = mask.unsqueeze(1)88 89        assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions"90        assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions"91        assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size"92 93        # Check image is in [-1, 1]94        if image.min() < -1 or image.max() > 1:95            raise ValueError("Image should be in [-1, 1] range")96 97        # Check mask is in [0, 1]98        if mask.min() < 0 or mask.max() > 1:99            raise ValueError("Mask should be in [0, 1] range")100 101        # Binarize mask102        mask[mask < 0.5] = 0103        mask[mask >= 0.5] = 1104 105        # Image as float32106        image = image.to(dtype=torch.float32)107    elif isinstance(mask, torch.Tensor):108        raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not")109    else:110        # preprocess image111        if isinstance(image, (PIL.Image.Image, np.ndarray)):112            image = [image]113 114        if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):115            image = [np.array(i.convert("RGB"))[None, :] for i in image]116            image = np.concatenate(image, axis=0)117        elif isinstance(image, list) and isinstance(image[0], np.ndarray):118            image = np.concatenate([i[None, :] for i in image], axis=0)119 120        image = image.transpose(0, 3, 1, 2)121        image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0122 123        # preprocess mask124        if isinstance(mask, (PIL.Image.Image, np.ndarray)):125            mask = [mask]126 127        if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):128            mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0)129            mask = mask.astype(np.float32) / 255.0130        elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):131            mask = np.concatenate([m[None, None, :] for m in mask], axis=0)132 133        mask[mask < 0.5] = 0134        mask[mask >= 0.5] = 1135        mask = torch.from_numpy(mask)136 137    # masked_image = image * (mask >= 0.5)138    masked_image = image139 140    return mask, masked_image141 142 143class StableDiffusionRepaintPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin):144    r"""145    Pipeline for text-guided image inpainting using Stable Diffusion. *This is an experimental feature*.146    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the147    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)148    In addition the pipeline inherits the following loading methods:149        - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`]150        - *LoRA*: [`loaders.LoraLoaderMixin.load_lora_weights`]151    as well as the following saving methods:152        - *LoRA*: [`loaders.LoraLoaderMixin.save_lora_weights`]153    Args:154        vae ([`AutoencoderKL`]):155            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.156        text_encoder ([`CLIPTextModel`]):157            Frozen text-encoder. Stable Diffusion uses the text portion of158            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically159            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.160        tokenizer (`CLIPTokenizer`):161            Tokenizer of class162            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).163        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.164        scheduler ([`SchedulerMixin`]):165            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of166            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].167        safety_checker ([`StableDiffusionSafetyChecker`]):168            Classification module that estimates whether generated images could be considered offensive or harmful.169            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.170        feature_extractor ([`CLIPImageProcessor`]):171            Model that extracts features from generated images to be used as inputs for the `safety_checker`.172    """173 174    _optional_components = ["safety_checker", "feature_extractor"]175 176    def __init__(177        self,178        vae: AutoencoderKL,179        text_encoder: CLIPTextModel,180        tokenizer: CLIPTokenizer,181        unet: UNet2DConditionModel,182        scheduler: KarrasDiffusionSchedulers,183        safety_checker: StableDiffusionSafetyChecker,184        feature_extractor: CLIPImageProcessor,185        requires_safety_checker: bool = True,186    ):187        super().__init__()188 189        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:190            deprecation_message = (191                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"192                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "193                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"194                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"195                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"196                " file"197            )198            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)199            new_config = dict(scheduler.config)200            new_config["steps_offset"] = 1201            scheduler._internal_dict = FrozenDict(new_config)202 203        if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:204            deprecation_message = (205                f"The configuration file of this scheduler: {scheduler} has not set the configuration"206                " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"207                " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"208                " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"209                " Hub, it would be very nice if you could open a Pull request for the"210                " `scheduler/scheduler_config.json` file"211            )212            deprecate(213                "skip_prk_steps not set",214                "1.0.0",215                deprecation_message,216                standard_warn=False,217            )218            new_config = dict(scheduler.config)219            new_config["skip_prk_steps"] = True220            scheduler._internal_dict = FrozenDict(new_config)221 222        if safety_checker is None and requires_safety_checker:223            logger.warning(224                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"225                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"226                " results in services or applications open to the public. Both the diffusers team and Hugging Face"227                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"228                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"229                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."230            )231 232        if safety_checker is not None and feature_extractor is None:233            raise ValueError(234                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"235                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."236            )237 238        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(239            version.parse(unet.config._diffusers_version).base_version240        ) < version.parse("0.9.0.dev0")241        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64242        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:243            deprecation_message = (244                "The configuration file of the unet has set the default `sample_size` to smaller than"245                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"246                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"247                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"248                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"249                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"250                " in the config might lead to incorrect results in future versions. If you have downloaded this"251                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"252                " the `unet/config.json` file"253            )254            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)255            new_config = dict(unet.config)256            new_config["sample_size"] = 64257            unet._internal_dict = FrozenDict(new_config)258        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4259        if unet.config.in_channels != 4:260            logger.warning(261                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"262                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"263                ". If you did not intend to modify"264                " this behavior, please check whether you have loaded the right checkpoint."265            )266 267        self.register_modules(268            vae=vae,269            text_encoder=text_encoder,270            tokenizer=tokenizer,271            unet=unet,272            scheduler=scheduler,273            safety_checker=safety_checker,274            feature_extractor=feature_extractor,275        )276        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)277        self.register_to_config(requires_safety_checker=requires_safety_checker)278 279    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_sequential_cpu_offload280    def enable_sequential_cpu_offload(self, gpu_id=0):281        r"""282        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,283        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a284        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.285        Note that offloading happens on a submodule basis. Memory savings are higher than with286        `enable_model_cpu_offload`, but performance is lower.287        """288        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):289            from accelerate import cpu_offload290        else:291            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")292 293        device = torch.device(f"cuda:{gpu_id}")294 295        if self.device.type != "cpu":296            self.to("cpu", silence_dtype_warnings=True)297            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)298 299        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:300            cpu_offload(cpu_offloaded_model, device)301 302        if self.safety_checker is not None:303            cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)304 305    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_model_cpu_offload306    def enable_model_cpu_offload(self, gpu_id=0):307        r"""308        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared309        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`310        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with311        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.312        """313        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):314            from accelerate import cpu_offload_with_hook315        else:316            raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")317 318        device = torch.device(f"cuda:{gpu_id}")319 320        if self.device.type != "cpu":321            self.to("cpu", silence_dtype_warnings=True)322            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)323 324        hook = None325        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:326            _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)327 328        if self.safety_checker is not None:329            _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)330 331        # We'll offload the last model manually.332        self.final_offload_hook = hook333 334    @property335    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device336    def _execution_device(self):337        r"""338        Returns the device on which the pipeline's models will be executed. After calling339        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module340        hooks.341        """342        if not hasattr(self.unet, "_hf_hook"):343            return self.device344        for module in self.unet.modules():345            if (346                hasattr(module, "_hf_hook")347                and hasattr(module._hf_hook, "execution_device")348                and module._hf_hook.execution_device is not None349            ):350                return torch.device(module._hf_hook.execution_device)351        return self.device352 353    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt354    def _encode_prompt(355        self,356        prompt,357        device,358        num_images_per_prompt,359        do_classifier_free_guidance,360        negative_prompt=None,361        prompt_embeds: Optional[torch.FloatTensor] = None,362        negative_prompt_embeds: Optional[torch.FloatTensor] = None,363    ):364        r"""365        Encodes the prompt into text encoder hidden states.366        Args:367             prompt (`str` or `List[str]`, *optional*):368                prompt to be encoded369            device: (`torch.device`):370                torch device371            num_images_per_prompt (`int`):372                number of images that should be generated per prompt373            do_classifier_free_guidance (`bool`):374                whether to use classifier free guidance or not375            negative_prompt (`str` or `List[str]`, *optional*):376                The prompt or prompts not to guide the image generation. If not defined, one has to pass377                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is378                less than `1`).379            prompt_embeds (`torch.FloatTensor`, *optional*):380                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not381                provided, text embeddings will be generated from `prompt` input argument.382            negative_prompt_embeds (`torch.FloatTensor`, *optional*):383                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt384                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input385                argument.386        """387        if prompt is not None and isinstance(prompt, str):388            batch_size = 1389        elif prompt is not None and isinstance(prompt, list):390            batch_size = len(prompt)391        else:392            batch_size = prompt_embeds.shape[0]393 394        if prompt_embeds is None:395            # textual inversion: procecss multi-vector tokens if necessary396            if isinstance(self, TextualInversionLoaderMixin):397                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)398 399            text_inputs = self.tokenizer(400                prompt,401                padding="max_length",402                max_length=self.tokenizer.model_max_length,403                truncation=True,404                return_tensors="pt",405            )406            text_input_ids = text_inputs.input_ids407            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids408 409            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(410                text_input_ids, untruncated_ids411            ):412                removed_text = self.tokenizer.batch_decode(413                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]414                )415                logger.warning(416                    "The following part of your input was truncated because CLIP can only handle sequences up to"417                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"418                )419 420            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:421                attention_mask = text_inputs.attention_mask.to(device)422            else:423                attention_mask = None424 425            prompt_embeds = self.text_encoder(426                text_input_ids.to(device),427                attention_mask=attention_mask,428            )429            prompt_embeds = prompt_embeds[0]430 431        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)432 433        bs_embed, seq_len, _ = prompt_embeds.shape434        # duplicate text embeddings for each generation per prompt, using mps friendly method435        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)436        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)437 438        # get unconditional embeddings for classifier free guidance439        if do_classifier_free_guidance and negative_prompt_embeds is None:440            uncond_tokens: List[str]441            if negative_prompt is None:442                uncond_tokens = [""] * batch_size443            elif type(prompt) is not type(negative_prompt):444                raise TypeError(445                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="446                    f" {type(prompt)}."447                )448            elif isinstance(negative_prompt, str):449                uncond_tokens = [negative_prompt]450            elif batch_size != len(negative_prompt):451                raise ValueError(452                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"453                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"454                    " the batch size of `prompt`."455                )456            else:457                uncond_tokens = negative_prompt458 459            # textual inversion: procecss multi-vector tokens if necessary460            if isinstance(self, TextualInversionLoaderMixin):461                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)462 463            max_length = prompt_embeds.shape[1]464            uncond_input = self.tokenizer(465                uncond_tokens,466                padding="max_length",467                max_length=max_length,468                truncation=True,469                return_tensors="pt",470            )471 472            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:473                attention_mask = uncond_input.attention_mask.to(device)474            else:475                attention_mask = None476 477            negative_prompt_embeds = self.text_encoder(478                uncond_input.input_ids.to(device),479                attention_mask=attention_mask,480            )481            negative_prompt_embeds = negative_prompt_embeds[0]482 483        if do_classifier_free_guidance:484            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method485            seq_len = negative_prompt_embeds.shape[1]486 487            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)488 489            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)490            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)491 492            # For classifier free guidance, we need to do two forward passes.493            # Here we concatenate the unconditional and text embeddings into a single batch494            # to avoid doing two forward passes495            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])496 497        return prompt_embeds498 499    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker500    def run_safety_checker(self, image, device, dtype):501        if self.safety_checker is not None:502            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)503            image, has_nsfw_concept = self.safety_checker(504                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)505            )506        else:507            has_nsfw_concept = None508        return image, has_nsfw_concept509 510    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs511    def prepare_extra_step_kwargs(self, generator, eta):512        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature513        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.514        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502515        # and should be between [0, 1]516 517        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())518        extra_step_kwargs = {}519        if accepts_eta:520            extra_step_kwargs["eta"] = eta521 522        # check if the scheduler accepts generator523        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())524        if accepts_generator:525            extra_step_kwargs["generator"] = generator526        return extra_step_kwargs527 528    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents529    def decode_latents(self, latents):530        latents = 1 / self.vae.config.scaling_factor * latents531        image = self.vae.decode(latents).sample532        image = (image / 2 + 0.5).clamp(0, 1)533        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16534        image = image.cpu().permute(0, 2, 3, 1).float().numpy()535        return image536 537    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs538    def check_inputs(539        self,540        prompt,541        height,542        width,543        callback_steps,544        negative_prompt=None,545        prompt_embeds=None,546        negative_prompt_embeds=None,547    ):548        if height % 8 != 0 or width % 8 != 0:549            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")550 551        if (callback_steps is None) or (552            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)553        ):554            raise ValueError(555                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"556                f" {type(callback_steps)}."557            )558 559        if prompt is not None and prompt_embeds is not None:560            raise ValueError(561                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"562                " only forward one of the two."563            )564        elif prompt is None and prompt_embeds is None:565            raise ValueError(566                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."567            )568        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):569            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")570 571        if negative_prompt is not None and negative_prompt_embeds is not None:572            raise ValueError(573                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"574                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."575            )576 577        if prompt_embeds is not None and negative_prompt_embeds is not None:578            if prompt_embeds.shape != negative_prompt_embeds.shape:579                raise ValueError(580                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"581                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"582                    f" {negative_prompt_embeds.shape}."583                )584 585    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents586    def prepare_latents(587        self,588        batch_size,589        num_channels_latents,590        height,591        width,592        dtype,593        device,594        generator,595        latents=None,596    ):597        shape = (598            batch_size,599            num_channels_latents,600            height // self.vae_scale_factor,601            width // self.vae_scale_factor,602        )603        if isinstance(generator, list) and len(generator) != batch_size:604            raise ValueError(605                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"606                f" size of {batch_size}. Make sure the batch size matches the length of the generators."607            )608 609        if latents is None:610            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)611        else:612            latents = latents.to(device)613 614        # scale the initial noise by the standard deviation required by the scheduler615        latents = latents * self.scheduler.init_noise_sigma616        return latents617 618    def prepare_mask_latents(619        self,620        mask,621        masked_image,622        batch_size,623        height,624        width,625        dtype,626        device,627        generator,628        do_classifier_free_guidance,629    ):630        # resize the mask to latents shape as we concatenate the mask to the latents631        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload632        # and half precision633        mask = torch.nn.functional.interpolate(634            mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)635        )636        mask = mask.to(device=device, dtype=dtype)637 638        masked_image = masked_image.to(device=device, dtype=dtype)639 640        # encode the mask image into latents space so we can concatenate it to the latents641        if isinstance(generator, list):642            masked_image_latents = [643                self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])644                for i in range(batch_size)645            ]646            masked_image_latents = torch.cat(masked_image_latents, dim=0)647        else:648            masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)649        masked_image_latents = self.vae.config.scaling_factor * masked_image_latents650 651        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method652        if mask.shape[0] < batch_size:653            if not batch_size % mask.shape[0] == 0:654                raise ValueError(655                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"656                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"657                    " of masks that you pass is divisible by the total requested batch size."658                )659            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)660        if masked_image_latents.shape[0] < batch_size:661            if not batch_size % masked_image_latents.shape[0] == 0:662                raise ValueError(663                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"664                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."665                    " Make sure the number of images that you pass is divisible by the total requested batch size."666                )667            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)668 669        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask670        masked_image_latents = (671            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents672        )673 674        # aligning device to prevent device errors when concating it with the latent model input675        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)676        return mask, masked_image_latents677 678    @torch.no_grad()679    def __call__(680        self,681        prompt: Union[str, List[str]] = None,682        image: Union[torch.FloatTensor, PIL.Image.Image] = None,683        mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,684        height: Optional[int] = None,685        width: Optional[int] = None,686        num_inference_steps: int = 50,687        jump_length: Optional[int] = 10,688        jump_n_sample: Optional[int] = 10,689        guidance_scale: float = 7.5,690        negative_prompt: Optional[Union[str, List[str]]] = None,691        num_images_per_prompt: Optional[int] = 1,692        eta: float = 0.0,693        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,694        latents: Optional[torch.FloatTensor] = None,695        prompt_embeds: Optional[torch.FloatTensor] = None,696        negative_prompt_embeds: Optional[torch.FloatTensor] = None,697        output_type: Optional[str] = "pil",698        return_dict: bool = True,699        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,700        callback_steps: int = 1,701    ):702        r"""703        Function invoked when calling the pipeline for generation.704        Args:705            prompt (`str` or `List[str]`, *optional*):706                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.707                instead.708            image (`PIL.Image.Image`):709                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will710                be masked out with `mask_image` and repainted according to `prompt`.711            mask_image (`PIL.Image.Image`):712                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be713                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted714                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)715                instead of 3, so the expected shape would be `(B, H, W, 1)`.716            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):717                The height in pixels of the generated image.718            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):719                The width in pixels of the generated image.720            num_inference_steps (`int`, *optional*, defaults to 50):721                The number of denoising steps. More denoising steps usually lead to a higher quality image at the722                expense of slower inference.723            jump_length (`int`, *optional*, defaults to 10):724                The number of steps taken forward in time before going backward in time for a single jump ("j" in725                RePaint paper). Take a look at Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.726            jump_n_sample (`int`, *optional*, defaults to 10):727                The number of times we will make forward time jump for a given chosen time sample. Take a look at728                Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.729            guidance_scale (`float`, *optional*, defaults to 7.5):730                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).731                `guidance_scale` is defined as `w` of equation 2. of [Imagen732                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >733                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,734                usually at the expense of lower image quality.735            negative_prompt (`str` or `List[str]`, *optional*):736                The prompt or prompts not to guide the image generation. If not defined, one has to pass737                `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`738                is less than `1`).739            num_images_per_prompt (`int`, *optional*, defaults to 1):740                The number of images to generate per prompt.741            eta (`float`, *optional*, defaults to 0.0):742                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to743                [`schedulers.DDIMScheduler`], will be ignored for others.744            generator (`torch.Generator`, *optional*):745                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)746                to make generation deterministic.747            latents (`torch.FloatTensor`, *optional*):748                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image749                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents750                tensor will ge generated by sampling using the supplied random `generator`.751            prompt_embeds (`torch.FloatTensor`, *optional*):752                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not753                provided, text embeddings will be generated from `prompt` input argument.754            negative_prompt_embeds (`torch.FloatTensor`, *optional*):755                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt756                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input757                argument.758            output_type (`str`, *optional*, defaults to `"pil"`):759                The output format of the generate image. Choose between760                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.761            return_dict (`bool`, *optional*, defaults to `True`):762                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a763                plain tuple.764            callback (`Callable`, *optional*):765                A function that will be called every `callback_steps` steps during inference. The function will be766                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.767            callback_steps (`int`, *optional*, defaults to 1):768                The frequency at which the `callback` function will be called. If not specified, the callback will be769                called at every step.770        Examples:771        ```py772        >>> import PIL773        >>> import requests774        >>> import torch775        >>> from io import BytesIO776        >>> from diffusers import StableDiffusionPipeline, RePaintScheduler777        >>> def download_image(url):778        ...     response = requests.get(url)779        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")780        >>> base_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/"781        >>> img_url = base_url + "overture-creations-5sI6fQgYIuo.png"782        >>> mask_url = base_url + "overture-creations-5sI6fQgYIuo_mask.png "783        >>> init_image = download_image(img_url).resize((512, 512))784        >>> mask_image = download_image(mask_url).resize((512, 512))785        >>> pipe = DiffusionPipeline.from_pretrained(786        ...     "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16, custom_pipeline="stable_diffusion_repaint",787        ... )788        >>> pipe.scheduler = RePaintScheduler.from_config(pipe.scheduler.config)789        >>> pipe = pipe.to("cuda")790        >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"791        >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]792        ```793        Returns:794            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:795            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.796            When returning a tuple, the first element is a list with the generated images, and the second element is a797            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"798            (nsfw) content, according to the `safety_checker`.799        """800        # 0. Default height and width to unet801        height = height or self.unet.config.sample_size * self.vae_scale_factor802        width = width or self.unet.config.sample_size * self.vae_scale_factor803 804        # 1. Check inputs805        self.check_inputs(806            prompt,807            height,808            width,809            callback_steps,810            negative_prompt,811            prompt_embeds,812            negative_prompt_embeds,813        )814 815        if image is None:816            raise ValueError("`image` input cannot be undefined.")817 818        if mask_image is None:819            raise ValueError("`mask_image` input cannot be undefined.")820 821        # 2. Define call parameters822        if prompt is not None and isinstance(prompt, str):823            batch_size = 1824        elif prompt is not None and isinstance(prompt, list):825            batch_size = len(prompt)826        else:827            batch_size = prompt_embeds.shape[0]828 829        device = self._execution_device830        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)831        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`832        # corresponds to doing no classifier free guidance.833        do_classifier_free_guidance = guidance_scale > 1.0834 835        # 3. Encode input prompt836        prompt_embeds = self._encode_prompt(837            prompt,838            device,839            num_images_per_prompt,840            do_classifier_free_guidance,841            negative_prompt,842            prompt_embeds=prompt_embeds,843            negative_prompt_embeds=negative_prompt_embeds,844        )845 846        # 4. Preprocess mask and image847        mask, masked_image = prepare_mask_and_masked_image(image, mask_image)848 849        # 5. set timesteps850        self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, device)851        self.scheduler.eta = eta852 853        timesteps = self.scheduler.timesteps854        # latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)855 856        # 6. Prepare latent variables857        num_channels_latents = self.vae.config.latent_channels858        latents = self.prepare_latents(859            batch_size * num_images_per_prompt,860            num_channels_latents,861            height,862            width,863            prompt_embeds.dtype,864            device,865            generator,866            latents,867        )868 869        # 7. Prepare mask latent variables870        mask, masked_image_latents = self.prepare_mask_latents(871            mask,872            masked_image,873            batch_size * num_images_per_prompt,874            height,875            width,876            prompt_embeds.dtype,877            device,878            generator,879            do_classifier_free_guidance=False,  # We do not need duplicate mask and image880        )881 882        # 8. Check that sizes of mask, masked image and latents match883        # num_channels_mask = mask.shape[1]884        # num_channels_masked_image = masked_image_latents.shape[1]885        if num_channels_latents != self.unet.config.in_channels:886            raise ValueError(887                f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"888                f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} "889                f" = Please verify the config of"890                " `pipeline.unet` or your `mask_image` or `image` input."891            )892 893        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline894        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)895 896        t_last = timesteps[0] + 1897 898        # 10. Denoising loop899        with self.progress_bar(total=len(timesteps)) as progress_bar:900            for i, t in enumerate(timesteps):901                if t >= t_last:902                    # compute the reverse: x_t-1 -> x_t903                    latents = self.scheduler.undo_step(latents, t_last, generator)904                    progress_bar.update()905                    t_last = t906                    continue907 908                # expand the latents if we are doing classifier free guidance909                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents910 911                # concat latents, mask, masked_image_latents in the channel dimension912                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)913                # latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)914 915                # predict the noise residual916                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample917 918                # perform guidance919                if do_classifier_free_guidance:920                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)921                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)922 923                # compute the previous noisy sample x_t -> x_t-1924                latents = self.scheduler.step(925                    noise_pred,926                    t,927                    latents,928                    masked_image_latents,929                    mask,930                    **extra_step_kwargs,931                ).prev_sample932 933                # call the callback, if provided934                progress_bar.update()935                if callback is not None and i % callback_steps == 0:936                    step_idx = i // getattr(self.scheduler, "order", 1)937                    callback(step_idx, t, latents)938 939                t_last = t940 941        # 11. Post-processing942        image = self.decode_latents(latents)943 944        # 12. Run safety checker945        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)946 947        # 13. Convert to PIL948        if output_type == "pil":949            image = self.numpy_to_pil(image)950 951        # Offload last model to CPU952        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:953            self.final_offload_hook.offload()954 955        if not return_dict:956            return (image, has_nsfw_concept)957 958        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)959