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    _optional_components = ["safety_checker", "feature_extractor"]174 175    def __init__(176        self,177        vae: AutoencoderKL,178        text_encoder: CLIPTextModel,179        tokenizer: CLIPTokenizer,180        unet: UNet2DConditionModel,181        scheduler: KarrasDiffusionSchedulers,182        safety_checker: StableDiffusionSafetyChecker,183        feature_extractor: CLIPImageProcessor,184        requires_safety_checker: bool = True,185    ):186        super().__init__()187 188        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:189            deprecation_message = (190                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"191                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "192                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"193                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"194                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"195                " file"196            )197            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)198            new_config = dict(scheduler.config)199            new_config["steps_offset"] = 1200            scheduler._internal_dict = FrozenDict(new_config)201 202        if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:203            deprecation_message = (204                f"The configuration file of this scheduler: {scheduler} has not set the configuration"205                " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"206                " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"207                " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"208                " Hub, it would be very nice if you could open a Pull request for the"209                " `scheduler/scheduler_config.json` file"210            )211            deprecate(212                "skip_prk_steps not set",213                "1.0.0",214                deprecation_message,215                standard_warn=False,216            )217            new_config = dict(scheduler.config)218            new_config["skip_prk_steps"] = True219            scheduler._internal_dict = FrozenDict(new_config)220 221        if safety_checker is None and requires_safety_checker:222            logger.warning(223                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"224                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"225                " results in services or applications open to the public. Both the diffusers team and Hugging Face"226                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"227                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"228                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."229            )230 231        if safety_checker is not None and feature_extractor is None:232            raise ValueError(233                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"234                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."235            )236 237        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(238            version.parse(unet.config._diffusers_version).base_version239        ) < version.parse("0.9.0.dev0")240        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64241        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:242            deprecation_message = (243                "The configuration file of the unet has set the default `sample_size` to smaller than"244                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"245                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"246                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"247                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"248                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"249                " in the config might lead to incorrect results in future versions. If you have downloaded this"250                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"251                " the `unet/config.json` file"252            )253            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)254            new_config = dict(unet.config)255            new_config["sample_size"] = 64256            unet._internal_dict = FrozenDict(new_config)257        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4258        if unet.config.in_channels != 4:259            logger.warning(260                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"261                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"262                ". If you did not intend to modify"263                " this behavior, please check whether you have loaded the right checkpoint."264            )265 266        self.register_modules(267            vae=vae,268            text_encoder=text_encoder,269            tokenizer=tokenizer,270            unet=unet,271            scheduler=scheduler,272            safety_checker=safety_checker,273            feature_extractor=feature_extractor,274        )275        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)276        self.register_to_config(requires_safety_checker=requires_safety_checker)277 278    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_sequential_cpu_offload279    def enable_sequential_cpu_offload(self, gpu_id=0):280        r"""281        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,282        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a283        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.284        Note that offloading happens on a submodule basis. Memory savings are higher than with285        `enable_model_cpu_offload`, but performance is lower.286        """287        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):288            from accelerate import cpu_offload289        else:290            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")291 292        device = torch.device(f"cuda:{gpu_id}")293 294        if self.device.type != "cpu":295            self.to("cpu", silence_dtype_warnings=True)296            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)297 298        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:299            cpu_offload(cpu_offloaded_model, device)300 301        if self.safety_checker is not None:302            cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)303 304    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_model_cpu_offload305    def enable_model_cpu_offload(self, gpu_id=0):306        r"""307        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared308        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`309        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with310        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.311        """312        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):313            from accelerate import cpu_offload_with_hook314        else:315            raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")316 317        device = torch.device(f"cuda:{gpu_id}")318 319        if self.device.type != "cpu":320            self.to("cpu", silence_dtype_warnings=True)321            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)322 323        hook = None324        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:325            _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)326 327        if self.safety_checker is not None:328            _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)329 330        # We'll offload the last model manually.331        self.final_offload_hook = hook332 333    @property334    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device335    def _execution_device(self):336        r"""337        Returns the device on which the pipeline's models will be executed. After calling338        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module339        hooks.340        """341        if not hasattr(self.unet, "_hf_hook"):342            return self.device343        for module in self.unet.modules():344            if (345                hasattr(module, "_hf_hook")346                and hasattr(module._hf_hook, "execution_device")347                and module._hf_hook.execution_device is not None348            ):349                return torch.device(module._hf_hook.execution_device)350        return self.device351 352    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt353    def _encode_prompt(354        self,355        prompt,356        device,357        num_images_per_prompt,358        do_classifier_free_guidance,359        negative_prompt=None,360        prompt_embeds: Optional[torch.FloatTensor] = None,361        negative_prompt_embeds: Optional[torch.FloatTensor] = None,362    ):363        r"""364        Encodes the prompt into text encoder hidden states.365        Args:366             prompt (`str` or `List[str]`, *optional*):367                prompt to be encoded368            device: (`torch.device`):369                torch device370            num_images_per_prompt (`int`):371                number of images that should be generated per prompt372            do_classifier_free_guidance (`bool`):373                whether to use classifier free guidance or not374            negative_prompt (`str` or `List[str]`, *optional*):375                The prompt or prompts not to guide the image generation. If not defined, one has to pass376                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is377                less than `1`).378            prompt_embeds (`torch.FloatTensor`, *optional*):379                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not380                provided, text embeddings will be generated from `prompt` input argument.381            negative_prompt_embeds (`torch.FloatTensor`, *optional*):382                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt383                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input384                argument.385        """386        if prompt is not None and isinstance(prompt, str):387            batch_size = 1388        elif prompt is not None and isinstance(prompt, list):389            batch_size = len(prompt)390        else:391            batch_size = prompt_embeds.shape[0]392 393        if prompt_embeds is None:394            # textual inversion: procecss multi-vector tokens if necessary395            if isinstance(self, TextualInversionLoaderMixin):396                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)397 398            text_inputs = self.tokenizer(399                prompt,400                padding="max_length",401                max_length=self.tokenizer.model_max_length,402                truncation=True,403                return_tensors="pt",404            )405            text_input_ids = text_inputs.input_ids406            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids407 408            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(409                text_input_ids, untruncated_ids410            ):411                removed_text = self.tokenizer.batch_decode(412                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]413                )414                logger.warning(415                    "The following part of your input was truncated because CLIP can only handle sequences up to"416                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"417                )418 419            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:420                attention_mask = text_inputs.attention_mask.to(device)421            else:422                attention_mask = None423 424            prompt_embeds = self.text_encoder(425                text_input_ids.to(device),426                attention_mask=attention_mask,427            )428            prompt_embeds = prompt_embeds[0]429 430        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)431 432        bs_embed, seq_len, _ = prompt_embeds.shape433        # duplicate text embeddings for each generation per prompt, using mps friendly method434        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)435        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)436 437        # get unconditional embeddings for classifier free guidance438        if do_classifier_free_guidance and negative_prompt_embeds is None:439            uncond_tokens: List[str]440            if negative_prompt is None:441                uncond_tokens = [""] * batch_size442            elif type(prompt) is not type(negative_prompt):443                raise TypeError(444                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="445                    f" {type(prompt)}."446                )447            elif isinstance(negative_prompt, str):448                uncond_tokens = [negative_prompt]449            elif batch_size != len(negative_prompt):450                raise ValueError(451                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"452                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"453                    " the batch size of `prompt`."454                )455            else:456                uncond_tokens = negative_prompt457 458            # textual inversion: procecss multi-vector tokens if necessary459            if isinstance(self, TextualInversionLoaderMixin):460                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)461 462            max_length = prompt_embeds.shape[1]463            uncond_input = self.tokenizer(464                uncond_tokens,465                padding="max_length",466                max_length=max_length,467                truncation=True,468                return_tensors="pt",469            )470 471            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:472                attention_mask = uncond_input.attention_mask.to(device)473            else:474                attention_mask = None475 476            negative_prompt_embeds = self.text_encoder(477                uncond_input.input_ids.to(device),478                attention_mask=attention_mask,479            )480            negative_prompt_embeds = negative_prompt_embeds[0]481 482        if do_classifier_free_guidance:483            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method484            seq_len = negative_prompt_embeds.shape[1]485 486            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)487 488            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)489            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)490 491            # For classifier free guidance, we need to do two forward passes.492            # Here we concatenate the unconditional and text embeddings into a single batch493            # to avoid doing two forward passes494            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])495 496        return prompt_embeds497 498    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker499    def run_safety_checker(self, image, device, dtype):500        if self.safety_checker is not None:501            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)502            image, has_nsfw_concept = self.safety_checker(503                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)504            )505        else:506            has_nsfw_concept = None507        return image, has_nsfw_concept508 509    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs510    def prepare_extra_step_kwargs(self, generator, eta):511        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature512        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.513        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502514        # and should be between [0, 1]515 516        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())517        extra_step_kwargs = {}518        if accepts_eta:519            extra_step_kwargs["eta"] = eta520 521        # check if the scheduler accepts generator522        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())523        if accepts_generator:524            extra_step_kwargs["generator"] = generator525        return extra_step_kwargs526 527    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents528    def decode_latents(self, latents):529        latents = 1 / self.vae.config.scaling_factor * latents530        image = self.vae.decode(latents).sample531        image = (image / 2 + 0.5).clamp(0, 1)532        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16533        image = image.cpu().permute(0, 2, 3, 1).float().numpy()534        return image535 536    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs537    def check_inputs(538        self,539        prompt,540        height,541        width,542        callback_steps,543        negative_prompt=None,544        prompt_embeds=None,545        negative_prompt_embeds=None,546    ):547        if height % 8 != 0 or width % 8 != 0:548            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")549 550        if (callback_steps is None) or (551            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)552        ):553            raise ValueError(554                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"555                f" {type(callback_steps)}."556            )557 558        if prompt is not None and prompt_embeds is not None:559            raise ValueError(560                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"561                " only forward one of the two."562            )563        elif prompt is None and prompt_embeds is None:564            raise ValueError(565                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."566            )567        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):568            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")569 570        if negative_prompt is not None and negative_prompt_embeds is not None:571            raise ValueError(572                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"573                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."574            )575 576        if prompt_embeds is not None and negative_prompt_embeds is not None:577            if prompt_embeds.shape != negative_prompt_embeds.shape:578                raise ValueError(579                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"580                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"581                    f" {negative_prompt_embeds.shape}."582                )583 584    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents585    def prepare_latents(586        self,587        batch_size,588        num_channels_latents,589        height,590        width,591        dtype,592        device,593        generator,594        latents=None,595    ):596        shape = (597            batch_size,598            num_channels_latents,599            height // self.vae_scale_factor,600            width // self.vae_scale_factor,601        )602        if isinstance(generator, list) and len(generator) != batch_size:603            raise ValueError(604                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"605                f" size of {batch_size}. Make sure the batch size matches the length of the generators."606            )607 608        if latents is None:609            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)610        else:611            latents = latents.to(device)612 613        # scale the initial noise by the standard deviation required by the scheduler614        latents = latents * self.scheduler.init_noise_sigma615        return latents616 617    def prepare_mask_latents(618        self,619        mask,620        masked_image,621        batch_size,622        height,623        width,624        dtype,625        device,626        generator,627        do_classifier_free_guidance,628    ):629        # resize the mask to latents shape as we concatenate the mask to the latents630        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload631        # and half precision632        mask = torch.nn.functional.interpolate(633            mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)634        )635        mask = mask.to(device=device, dtype=dtype)636 637        masked_image = masked_image.to(device=device, dtype=dtype)638 639        # encode the mask image into latents space so we can concatenate it to the latents640        if isinstance(generator, list):641            masked_image_latents = [642                self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])643                for i in range(batch_size)644            ]645            masked_image_latents = torch.cat(masked_image_latents, dim=0)646        else:647            masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)648        masked_image_latents = self.vae.config.scaling_factor * masked_image_latents649 650        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method651        if mask.shape[0] < batch_size:652            if not batch_size % mask.shape[0] == 0:653                raise ValueError(654                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"655                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"656                    " of masks that you pass is divisible by the total requested batch size."657                )658            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)659        if masked_image_latents.shape[0] < batch_size:660            if not batch_size % masked_image_latents.shape[0] == 0:661                raise ValueError(662                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"663                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."664                    " Make sure the number of images that you pass is divisible by the total requested batch size."665                )666            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)667 668        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask669        masked_image_latents = (670            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents671        )672 673        # aligning device to prevent device errors when concating it with the latent model input674        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)675        return mask, masked_image_latents676 677    @torch.no_grad()678    def __call__(679        self,680        prompt: Union[str, List[str]] = None,681        image: Union[torch.FloatTensor, PIL.Image.Image] = None,682        mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,683        height: Optional[int] = None,684        width: Optional[int] = None,685        num_inference_steps: int = 50,686        jump_length: Optional[int] = 10,687        jump_n_sample: Optional[int] = 10,688        guidance_scale: float = 7.5,689        negative_prompt: Optional[Union[str, List[str]]] = None,690        num_images_per_prompt: Optional[int] = 1,691        eta: float = 0.0,692        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,693        latents: Optional[torch.FloatTensor] = None,694        prompt_embeds: Optional[torch.FloatTensor] = None,695        negative_prompt_embeds: Optional[torch.FloatTensor] = None,696        output_type: Optional[str] = "pil",697        return_dict: bool = True,698        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,699        callback_steps: int = 1,700    ):701        r"""702        Function invoked when calling the pipeline for generation.703        Args:704            prompt (`str` or `List[str]`, *optional*):705                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.706                instead.707            image (`PIL.Image.Image`):708                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will709                be masked out with `mask_image` and repainted according to `prompt`.710            mask_image (`PIL.Image.Image`):711                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be712                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted713                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)714                instead of 3, so the expected shape would be `(B, H, W, 1)`.715            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):716                The height in pixels of the generated image.717            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):718                The width in pixels of the generated image.719            num_inference_steps (`int`, *optional*, defaults to 50):720                The number of denoising steps. More denoising steps usually lead to a higher quality image at the721                expense of slower inference.722            jump_length (`int`, *optional*, defaults to 10):723                The number of steps taken forward in time before going backward in time for a single jump ("j" in724                RePaint paper). Take a look at Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.725            jump_n_sample (`int`, *optional*, defaults to 10):726                The number of times we will make forward time jump for a given chosen time sample. Take a look at727                Figure 9 and 10 in https://arxiv.org/pdf/2201.09865.pdf.728            guidance_scale (`float`, *optional*, defaults to 7.5):729                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).730                `guidance_scale` is defined as `w` of equation 2. of [Imagen731                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >732                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,733                usually at the expense of lower image quality.734            negative_prompt (`str` or `List[str]`, *optional*):735                The prompt or prompts not to guide the image generation. If not defined, one has to pass736                `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale`737                is less than `1`).738            num_images_per_prompt (`int`, *optional*, defaults to 1):739                The number of images to generate per prompt.740            eta (`float`, *optional*, defaults to 0.0):741                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to742                [`schedulers.DDIMScheduler`], will be ignored for others.743            generator (`torch.Generator`, *optional*):744                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)745                to make generation deterministic.746            latents (`torch.FloatTensor`, *optional*):747                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image748                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents749                tensor will ge generated by sampling using the supplied random `generator`.750            prompt_embeds (`torch.FloatTensor`, *optional*):751                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not752                provided, text embeddings will be generated from `prompt` input argument.753            negative_prompt_embeds (`torch.FloatTensor`, *optional*):754                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt755                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input756                argument.757            output_type (`str`, *optional*, defaults to `"pil"`):758                The output format of the generate image. Choose between759                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.760            return_dict (`bool`, *optional*, defaults to `True`):761                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a762                plain tuple.763            callback (`Callable`, *optional*):764                A function that will be called every `callback_steps` steps during inference. The function will be765                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.766            callback_steps (`int`, *optional*, defaults to 1):767                The frequency at which the `callback` function will be called. If not specified, the callback will be768                called at every step.769        Examples:770        ```py771        >>> import PIL772        >>> import requests773        >>> import torch774        >>> from io import BytesIO775        >>> from diffusers import StableDiffusionPipeline, RePaintScheduler776        >>> def download_image(url):777        ...     response = requests.get(url)778        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")779        >>> base_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/"780        >>> img_url = base_url + "overture-creations-5sI6fQgYIuo.png"781        >>> mask_url = base_url + "overture-creations-5sI6fQgYIuo_mask.png "782        >>> init_image = download_image(img_url).resize((512, 512))783        >>> mask_image = download_image(mask_url).resize((512, 512))784        >>> pipe = DiffusionPipeline.from_pretrained(785        ...     "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16, custom_pipeline="stable_diffusion_repaint",786        ... )787        >>> pipe.scheduler = RePaintScheduler.from_config(pipe.scheduler.config)788        >>> pipe = pipe.to("cuda")789        >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"790        >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]791        ```792        Returns:793            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:794            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.795            When returning a tuple, the first element is a list with the generated images, and the second element is a796            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"797            (nsfw) content, according to the `safety_checker`.798        """799        # 0. Default height and width to unet800        height = height or self.unet.config.sample_size * self.vae_scale_factor801        width = width or self.unet.config.sample_size * self.vae_scale_factor802 803        # 1. Check inputs804        self.check_inputs(805            prompt,806            height,807            width,808            callback_steps,809            negative_prompt,810            prompt_embeds,811            negative_prompt_embeds,812        )813 814        if image is None:815            raise ValueError("`image` input cannot be undefined.")816 817        if mask_image is None:818            raise ValueError("`mask_image` input cannot be undefined.")819 820        # 2. Define call parameters821        if prompt is not None and isinstance(prompt, str):822            batch_size = 1823        elif prompt is not None and isinstance(prompt, list):824            batch_size = len(prompt)825        else:826            batch_size = prompt_embeds.shape[0]827 828        device = self._execution_device829        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)830        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`831        # corresponds to doing no classifier free guidance.832        do_classifier_free_guidance = guidance_scale > 1.0833 834        # 3. Encode input prompt835        prompt_embeds = self._encode_prompt(836            prompt,837            device,838            num_images_per_prompt,839            do_classifier_free_guidance,840            negative_prompt,841            prompt_embeds=prompt_embeds,842            negative_prompt_embeds=negative_prompt_embeds,843        )844 845        # 4. Preprocess mask and image846        mask, masked_image = prepare_mask_and_masked_image(image, mask_image)847 848        # 5. set timesteps849        self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, device)850        self.scheduler.eta = eta851 852        timesteps = self.scheduler.timesteps853        # latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)854 855        # 6. Prepare latent variables856        num_channels_latents = self.vae.config.latent_channels857        latents = self.prepare_latents(858            batch_size * num_images_per_prompt,859            num_channels_latents,860            height,861            width,862            prompt_embeds.dtype,863            device,864            generator,865            latents,866        )867 868        # 7. Prepare mask latent variables869        mask, masked_image_latents = self.prepare_mask_latents(870            mask,871            masked_image,872            batch_size * num_images_per_prompt,873            height,874            width,875            prompt_embeds.dtype,876            device,877            generator,878            do_classifier_free_guidance=False,  # We do not need duplicate mask and image879        )880 881        # 8. Check that sizes of mask, masked image and latents match882        # num_channels_mask = mask.shape[1]883        # num_channels_masked_image = masked_image_latents.shape[1]884        if num_channels_latents != self.unet.config.in_channels:885            raise ValueError(886                f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"887                f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} "888                f" = Please verify the config of"889                " `pipeline.unet` or your `mask_image` or `image` input."890            )891 892        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline893        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)894 895        t_last = timesteps[0] + 1896 897        # 10. Denoising loop898        with self.progress_bar(total=len(timesteps)) as progress_bar:899            for i, t in enumerate(timesteps):900                if t >= t_last:901                    # compute the reverse: x_t-1 -> x_t902                    latents = self.scheduler.undo_step(latents, t_last, generator)903                    progress_bar.update()904                    t_last = t905                    continue906 907                # expand the latents if we are doing classifier free guidance908                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents909 910                # concat latents, mask, masked_image_latents in the channel dimension911                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)912                # latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)913 914                # predict the noise residual915                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample916 917                # perform guidance918                if do_classifier_free_guidance:919                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)920                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)921 922                # compute the previous noisy sample x_t -> x_t-1923                latents = self.scheduler.step(924                    noise_pred,925                    t,926                    latents,927                    masked_image_latents,928                    mask,929                    **extra_step_kwargs,930                ).prev_sample931 932                # call the callback, if provided933                progress_bar.update()934                if callback is not None and i % callback_steps == 0:935                    step_idx = i // getattr(self.scheduler, "order", 1)936                    callback(step_idx, t, latents)937 938                t_last = t939 940        # 11. Post-processing941        image = self.decode_latents(latents)942 943        # 12. Run safety checker944        image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)945 946        # 13. Convert to PIL947        if output_type == "pil":948            image = self.numpy_to_pil(image)949 950        # Offload last model to CPU951        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:952            self.final_offload_hook.offload()953 954        if not return_dict:955            return (image, has_nsfw_concept)956 957        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)958