diffusers/community-pipelines-mirror

1# Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/2 3import inspect4from typing import Any, Callable, Dict, List, Optional, Tuple, Union5 6import numpy as np7import PIL.Image8import torch9import torch.nn.functional as F10from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer11 12from diffusers import AutoencoderKL, ControlNetModel, UNet2DConditionModel, logging13from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel14from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin15from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker16from diffusers.schedulers import KarrasDiffusionSchedulers17from diffusers.utils import (18    PIL_INTERPOLATION,19    replace_example_docstring,20)21from diffusers.utils.torch_utils import randn_tensor22 23 24logger = logging.get_logger(__name__)  # pylint: disable=invalid-name25 26EXAMPLE_DOC_STRING = """27    Examples:28        ```py29        >>> import numpy as np30        >>> import torch31        >>> from PIL import Image32        >>> from stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline33 34        >>> from transformers import AutoImageProcessor, UperNetForSemanticSegmentation35        >>> from diffusers import ControlNetModel, UniPCMultistepScheduler36        >>> from diffusers.utils import load_image37 38        >>> def ade_palette():39                return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],40                        [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],41                        [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],42                        [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],43                        [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],44                        [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],45                        [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],46                        [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],47                        [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],48                        [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],49                        [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],50                        [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],51                        [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],52                        [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],53                        [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],54                        [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],55                        [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],56                        [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],57                        [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],58                        [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],59                        [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],60                        [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],61                        [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],62                        [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],63                        [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],64                        [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],65                        [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],66                        [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],67                        [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],68                        [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],69                        [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],70                        [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],71                        [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],72                        [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],73                        [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],74                        [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],75                        [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],76                        [102, 255, 0], [92, 0, 255]]77 78        >>> image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")79        >>> image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")80 81        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16)82 83        >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(84                "runwayml/stable-diffusion-inpainting", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float1685            )86 87        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)88        >>> pipe.enable_xformers_memory_efficient_attention()89        >>> pipe.enable_model_cpu_offload()90 91        >>> def image_to_seg(image):92                pixel_values = image_processor(image, return_tensors="pt").pixel_values93                with torch.no_grad():94                    outputs = image_segmentor(pixel_values)95                seg = image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]96                color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)  # height, width, 397                palette = np.array(ade_palette())98                for label, color in enumerate(palette):99                    color_seg[seg == label, :] = color100                color_seg = color_seg.astype(np.uint8)101                seg_image = Image.fromarray(color_seg)102                return seg_image103 104        >>> image = load_image(105                "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"106            )107 108        >>> mask_image = load_image(109                "https://github.com/CompVis/latent-diffusion/raw/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"110            )111 112        >>> controlnet_conditioning_image = image_to_seg(image)113 114        >>> image = pipe(115                "Face of a yellow cat, high resolution, sitting on a park bench",116                image,117                mask_image,118                controlnet_conditioning_image,119                num_inference_steps=20,120            ).images[0]121 122        >>> image.save("out.png")123        ```124"""125 126 127def prepare_image(image):128    if isinstance(image, torch.Tensor):129        # Batch single image130        if image.ndim == 3:131            image = image.unsqueeze(0)132 133        image = image.to(dtype=torch.float32)134    else:135        # preprocess image136        if isinstance(image, (PIL.Image.Image, np.ndarray)):137            image = [image]138 139        if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):140            image = [np.array(i.convert("RGB"))[None, :] for i in image]141            image = np.concatenate(image, axis=0)142        elif isinstance(image, list) and isinstance(image[0], np.ndarray):143            image = np.concatenate([i[None, :] for i in image], axis=0)144 145        image = image.transpose(0, 3, 1, 2)146        image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0147 148    return image149 150 151def prepare_mask_image(mask_image):152    if isinstance(mask_image, torch.Tensor):153        if mask_image.ndim == 2:154            # Batch and add channel dim for single mask155            mask_image = mask_image.unsqueeze(0).unsqueeze(0)156        elif mask_image.ndim == 3 and mask_image.shape[0] == 1:157            # Single mask, the 0'th dimension is considered to be158            # the existing batch size of 1159            mask_image = mask_image.unsqueeze(0)160        elif mask_image.ndim == 3 and mask_image.shape[0] != 1:161            # Batch of mask, the 0'th dimension is considered to be162            # the batching dimension163            mask_image = mask_image.unsqueeze(1)164 165        # Binarize mask166        mask_image[mask_image < 0.5] = 0167        mask_image[mask_image >= 0.5] = 1168    else:169        # preprocess mask170        if isinstance(mask_image, (PIL.Image.Image, np.ndarray)):171            mask_image = [mask_image]172 173        if isinstance(mask_image, list) and isinstance(mask_image[0], PIL.Image.Image):174            mask_image = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask_image], axis=0)175            mask_image = mask_image.astype(np.float32) / 255.0176        elif isinstance(mask_image, list) and isinstance(mask_image[0], np.ndarray):177            mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0)178 179        mask_image[mask_image < 0.5] = 0180        mask_image[mask_image >= 0.5] = 1181        mask_image = torch.from_numpy(mask_image)182 183    return mask_image184 185 186def prepare_controlnet_conditioning_image(187    controlnet_conditioning_image,188    width,189    height,190    batch_size,191    num_images_per_prompt,192    device,193    dtype,194    do_classifier_free_guidance,195):196    if not isinstance(controlnet_conditioning_image, torch.Tensor):197        if isinstance(controlnet_conditioning_image, PIL.Image.Image):198            controlnet_conditioning_image = [controlnet_conditioning_image]199 200        if isinstance(controlnet_conditioning_image[0], PIL.Image.Image):201            controlnet_conditioning_image = [202                np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :]203                for i in controlnet_conditioning_image204            ]205            controlnet_conditioning_image = np.concatenate(controlnet_conditioning_image, axis=0)206            controlnet_conditioning_image = np.array(controlnet_conditioning_image).astype(np.float32) / 255.0207            controlnet_conditioning_image = controlnet_conditioning_image.transpose(0, 3, 1, 2)208            controlnet_conditioning_image = torch.from_numpy(controlnet_conditioning_image)209        elif isinstance(controlnet_conditioning_image[0], torch.Tensor):210            controlnet_conditioning_image = torch.cat(controlnet_conditioning_image, dim=0)211 212    image_batch_size = controlnet_conditioning_image.shape[0]213 214    if image_batch_size == 1:215        repeat_by = batch_size216    else:217        # image batch size is the same as prompt batch size218        repeat_by = num_images_per_prompt219 220    controlnet_conditioning_image = controlnet_conditioning_image.repeat_interleave(repeat_by, dim=0)221 222    controlnet_conditioning_image = controlnet_conditioning_image.to(device=device, dtype=dtype)223 224    if do_classifier_free_guidance:225        controlnet_conditioning_image = torch.cat([controlnet_conditioning_image] * 2)226 227    return controlnet_conditioning_image228 229 230class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline, StableDiffusionMixin):231    """232    Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/233    """234 235    _optional_components = ["safety_checker", "feature_extractor"]236 237    def __init__(238        self,239        vae: AutoencoderKL,240        text_encoder: CLIPTextModel,241        tokenizer: CLIPTokenizer,242        unet: UNet2DConditionModel,243        controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],244        scheduler: KarrasDiffusionSchedulers,245        safety_checker: StableDiffusionSafetyChecker,246        feature_extractor: CLIPImageProcessor,247        requires_safety_checker: bool = True,248    ):249        super().__init__()250 251        if safety_checker is None and requires_safety_checker:252            logger.warning(253                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"254                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"255                " results in services or applications open to the public. Both the diffusers team and Hugging Face"256                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"257                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"258                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."259            )260 261        if safety_checker is not None and feature_extractor is None:262            raise ValueError(263                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"264                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."265            )266 267        if isinstance(controlnet, (list, tuple)):268            controlnet = MultiControlNetModel(controlnet)269 270        self.register_modules(271            vae=vae,272            text_encoder=text_encoder,273            tokenizer=tokenizer,274            unet=unet,275            controlnet=controlnet,276            scheduler=scheduler,277            safety_checker=safety_checker,278            feature_extractor=feature_extractor,279        )280 281        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)282        self.register_to_config(requires_safety_checker=requires_safety_checker)283 284    def _encode_prompt(285        self,286        prompt,287        device,288        num_images_per_prompt,289        do_classifier_free_guidance,290        negative_prompt=None,291        prompt_embeds: Optional[torch.Tensor] = None,292        negative_prompt_embeds: Optional[torch.Tensor] = None,293    ):294        r"""295        Encodes the prompt into text encoder hidden states.296 297        Args:298             prompt (`str` or `List[str]`, *optional*):299                prompt to be encoded300            device: (`torch.device`):301                torch device302            num_images_per_prompt (`int`):303                number of images that should be generated per prompt304            do_classifier_free_guidance (`bool`):305                whether to use classifier free guidance or not306            negative_prompt (`str` or `List[str]`, *optional*):307                The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.308                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).309            prompt_embeds (`torch.Tensor`, *optional*):310                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not311                provided, text embeddings will be generated from `prompt` input argument.312            negative_prompt_embeds (`torch.Tensor`, *optional*):313                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt314                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input315                argument.316        """317        if prompt is not None and isinstance(prompt, str):318            batch_size = 1319        elif prompt is not None and isinstance(prompt, list):320            batch_size = len(prompt)321        else:322            batch_size = prompt_embeds.shape[0]323 324        if prompt_embeds is None:325            text_inputs = self.tokenizer(326                prompt,327                padding="max_length",328                max_length=self.tokenizer.model_max_length,329                truncation=True,330                return_tensors="pt",331            )332            text_input_ids = text_inputs.input_ids333            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids334 335            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(336                text_input_ids, untruncated_ids337            ):338                removed_text = self.tokenizer.batch_decode(339                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]340                )341                logger.warning(342                    "The following part of your input was truncated because CLIP can only handle sequences up to"343                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"344                )345 346            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:347                attention_mask = text_inputs.attention_mask.to(device)348            else:349                attention_mask = None350 351            prompt_embeds = self.text_encoder(352                text_input_ids.to(device),353                attention_mask=attention_mask,354            )355            prompt_embeds = prompt_embeds[0]356 357        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)358 359        bs_embed, seq_len, _ = prompt_embeds.shape360        # duplicate text embeddings for each generation per prompt, using mps friendly method361        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)362        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)363 364        # get unconditional embeddings for classifier free guidance365        if do_classifier_free_guidance and negative_prompt_embeds is None:366            uncond_tokens: List[str]367            if negative_prompt is None:368                uncond_tokens = [""] * batch_size369            elif type(prompt) is not type(negative_prompt):370                raise TypeError(371                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="372                    f" {type(prompt)}."373                )374            elif isinstance(negative_prompt, str):375                uncond_tokens = [negative_prompt]376            elif batch_size != len(negative_prompt):377                raise ValueError(378                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"379                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"380                    " the batch size of `prompt`."381                )382            else:383                uncond_tokens = negative_prompt384 385            max_length = prompt_embeds.shape[1]386            uncond_input = self.tokenizer(387                uncond_tokens,388                padding="max_length",389                max_length=max_length,390                truncation=True,391                return_tensors="pt",392            )393 394            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:395                attention_mask = uncond_input.attention_mask.to(device)396            else:397                attention_mask = None398 399            negative_prompt_embeds = self.text_encoder(400                uncond_input.input_ids.to(device),401                attention_mask=attention_mask,402            )403            negative_prompt_embeds = negative_prompt_embeds[0]404 405        if do_classifier_free_guidance:406            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method407            seq_len = negative_prompt_embeds.shape[1]408 409            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)410 411            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)412            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)413 414            # For classifier free guidance, we need to do two forward passes.415            # Here we concatenate the unconditional and text embeddings into a single batch416            # to avoid doing two forward passes417            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])418 419        return prompt_embeds420 421    def run_safety_checker(self, image, device, dtype):422        if self.safety_checker is not None:423            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)424            image, has_nsfw_concept = self.safety_checker(425                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)426            )427        else:428            has_nsfw_concept = None429        return image, has_nsfw_concept430 431    def decode_latents(self, latents):432        latents = 1 / self.vae.config.scaling_factor * latents433        image = self.vae.decode(latents).sample434        image = (image / 2 + 0.5).clamp(0, 1)435        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16436        image = image.cpu().permute(0, 2, 3, 1).float().numpy()437        return image438 439    def prepare_extra_step_kwargs(self, generator, eta):440        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature441        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.442        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502443        # and should be between [0, 1]444 445        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())446        extra_step_kwargs = {}447        if accepts_eta:448            extra_step_kwargs["eta"] = eta449 450        # check if the scheduler accepts generator451        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())452        if accepts_generator:453            extra_step_kwargs["generator"] = generator454        return extra_step_kwargs455 456    def check_controlnet_conditioning_image(self, image, prompt, prompt_embeds):457        image_is_pil = isinstance(image, PIL.Image.Image)458        image_is_tensor = isinstance(image, torch.Tensor)459        image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)460        image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)461 462        if not image_is_pil and not image_is_tensor and not image_is_pil_list and not image_is_tensor_list:463            raise TypeError(464                "image must be passed and be one of PIL image, torch tensor, list of PIL images, or list of torch tensors"465            )466 467        if image_is_pil:468            image_batch_size = 1469        elif image_is_tensor:470            image_batch_size = image.shape[0]471        elif image_is_pil_list:472            image_batch_size = len(image)473        elif image_is_tensor_list:474            image_batch_size = len(image)475        else:476            raise ValueError("controlnet condition image is not valid")477 478        if prompt is not None and isinstance(prompt, str):479            prompt_batch_size = 1480        elif prompt is not None and isinstance(prompt, list):481            prompt_batch_size = len(prompt)482        elif prompt_embeds is not None:483            prompt_batch_size = prompt_embeds.shape[0]484        else:485            raise ValueError("prompt or prompt_embeds are not valid")486 487        if image_batch_size != 1 and image_batch_size != prompt_batch_size:488            raise ValueError(489                f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"490            )491 492    def check_inputs(493        self,494        prompt,495        image,496        mask_image,497        controlnet_conditioning_image,498        height,499        width,500        callback_steps,501        negative_prompt=None,502        prompt_embeds=None,503        negative_prompt_embeds=None,504        controlnet_conditioning_scale=None,505    ):506        if height % 8 != 0 or width % 8 != 0:507            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")508 509        if (callback_steps is None) or (510            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)511        ):512            raise ValueError(513                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"514                f" {type(callback_steps)}."515            )516 517        if prompt is not None and prompt_embeds is not None:518            raise ValueError(519                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"520                " only forward one of the two."521            )522        elif prompt is None and prompt_embeds is None:523            raise ValueError(524                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."525            )526        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):527            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")528 529        if negative_prompt is not None and negative_prompt_embeds is not None:530            raise ValueError(531                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"532                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."533            )534 535        if prompt_embeds is not None and negative_prompt_embeds is not None:536            if prompt_embeds.shape != negative_prompt_embeds.shape:537                raise ValueError(538                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"539                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"540                    f" {negative_prompt_embeds.shape}."541                )542 543        # check controlnet condition image544        if isinstance(self.controlnet, ControlNetModel):545            self.check_controlnet_conditioning_image(controlnet_conditioning_image, prompt, prompt_embeds)546        elif isinstance(self.controlnet, MultiControlNetModel):547            if not isinstance(controlnet_conditioning_image, list):548                raise TypeError("For multiple controlnets: `image` must be type `list`")549            if len(controlnet_conditioning_image) != len(self.controlnet.nets):550                raise ValueError(551                    "For multiple controlnets: `image` must have the same length as the number of controlnets."552                )553            for image_ in controlnet_conditioning_image:554                self.check_controlnet_conditioning_image(image_, prompt, prompt_embeds)555        else:556            assert False557 558        # Check `controlnet_conditioning_scale`559        if isinstance(self.controlnet, ControlNetModel):560            if not isinstance(controlnet_conditioning_scale, float):561                raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")562        elif isinstance(self.controlnet, MultiControlNetModel):563            if isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(564                self.controlnet.nets565            ):566                raise ValueError(567                    "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"568                    " the same length as the number of controlnets"569                )570        else:571            assert False572 573        if isinstance(image, torch.Tensor) and not isinstance(mask_image, torch.Tensor):574            raise TypeError("if `image` is a tensor, `mask_image` must also be a tensor")575 576        if isinstance(image, PIL.Image.Image) and not isinstance(mask_image, PIL.Image.Image):577            raise TypeError("if `image` is a PIL image, `mask_image` must also be a PIL image")578 579        if isinstance(image, torch.Tensor):580            if image.ndim != 3 and image.ndim != 4:581                raise ValueError("`image` must have 3 or 4 dimensions")582 583            if mask_image.ndim != 2 and mask_image.ndim != 3 and mask_image.ndim != 4:584                raise ValueError("`mask_image` must have 2, 3, or 4 dimensions")585 586            if image.ndim == 3:587                image_batch_size = 1588                image_channels, image_height, image_width = image.shape589            elif image.ndim == 4:590                image_batch_size, image_channels, image_height, image_width = image.shape591            else:592                assert False593 594            if mask_image.ndim == 2:595                mask_image_batch_size = 1596                mask_image_channels = 1597                mask_image_height, mask_image_width = mask_image.shape598            elif mask_image.ndim == 3:599                mask_image_channels = 1600                mask_image_batch_size, mask_image_height, mask_image_width = mask_image.shape601            elif mask_image.ndim == 4:602                mask_image_batch_size, mask_image_channels, mask_image_height, mask_image_width = mask_image.shape603 604            if image_channels != 3:605                raise ValueError("`image` must have 3 channels")606 607            if mask_image_channels != 1:608                raise ValueError("`mask_image` must have 1 channel")609 610            if image_batch_size != mask_image_batch_size:611                raise ValueError("`image` and `mask_image` mush have the same batch sizes")612 613            if image_height != mask_image_height or image_width != mask_image_width:614                raise ValueError("`image` and `mask_image` must have the same height and width dimensions")615 616            if image.min() < -1 or image.max() > 1:617                raise ValueError("`image` should be in range [-1, 1]")618 619            if mask_image.min() < 0 or mask_image.max() > 1:620                raise ValueError("`mask_image` should be in range [0, 1]")621        else:622            mask_image_channels = 1623            image_channels = 3624 625        single_image_latent_channels = self.vae.config.latent_channels626 627        total_latent_channels = single_image_latent_channels * 2 + mask_image_channels628 629        if total_latent_channels != self.unet.config.in_channels:630            raise ValueError(631                f"The config of `pipeline.unet` expects {self.unet.config.in_channels} but received"632                f" non inpainting latent channels: {single_image_latent_channels},"633                f" mask channels: {mask_image_channels}, and masked image channels: {single_image_latent_channels}."634                f" Please verify the config of `pipeline.unet` and the `mask_image` and `image` inputs."635            )636 637    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):638        shape = (639            batch_size,640            num_channels_latents,641            int(height) // self.vae_scale_factor,642            int(width) // self.vae_scale_factor,643        )644        if isinstance(generator, list) and len(generator) != batch_size:645            raise ValueError(646                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"647                f" size of {batch_size}. Make sure the batch size matches the length of the generators."648            )649 650        if latents is None:651            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)652        else:653            latents = latents.to(device)654 655        # scale the initial noise by the standard deviation required by the scheduler656        latents = latents * self.scheduler.init_noise_sigma657 658        return latents659 660    def prepare_mask_latents(self, mask_image, batch_size, height, width, dtype, device, do_classifier_free_guidance):661        # resize the mask to latents shape as we concatenate the mask to the latents662        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload663        # and half precision664        mask_image = F.interpolate(mask_image, size=(height // self.vae_scale_factor, width // self.vae_scale_factor))665        mask_image = mask_image.to(device=device, dtype=dtype)666 667        # duplicate mask for each generation per prompt, using mps friendly method668        if mask_image.shape[0] < batch_size:669            if not batch_size % mask_image.shape[0] == 0:670                raise ValueError(671                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"672                    f" a total batch size of {batch_size}, but {mask_image.shape[0]} masks were passed. Make sure the number"673                    " of masks that you pass is divisible by the total requested batch size."674                )675            mask_image = mask_image.repeat(batch_size // mask_image.shape[0], 1, 1, 1)676 677        mask_image = torch.cat([mask_image] * 2) if do_classifier_free_guidance else mask_image678 679        mask_image_latents = mask_image680 681        return mask_image_latents682 683    def prepare_masked_image_latents(684        self, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance685    ):686        masked_image = masked_image.to(device=device, dtype=dtype)687 688        # encode the mask image into latents space so we can concatenate it to the latents689        if isinstance(generator, list):690            masked_image_latents = [691                self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i])692                for i in range(batch_size)693            ]694            masked_image_latents = torch.cat(masked_image_latents, dim=0)695        else:696            masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)697        masked_image_latents = self.vae.config.scaling_factor * masked_image_latents698 699        # duplicate masked_image_latents for each generation per prompt, using mps friendly method700        if masked_image_latents.shape[0] < batch_size:701            if not batch_size % masked_image_latents.shape[0] == 0:702                raise ValueError(703                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"704                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."705                    " Make sure the number of images that you pass is divisible by the total requested batch size."706                )707            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)708 709        masked_image_latents = (710            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents711        )712 713        # aligning device to prevent device errors when concating it with the latent model input714        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)715        return masked_image_latents716 717    def _default_height_width(self, height, width, image):718        if isinstance(image, list):719            image = image[0]720 721        if height is None:722            if isinstance(image, PIL.Image.Image):723                height = image.height724            elif isinstance(image, torch.Tensor):725                height = image.shape[3]726 727            height = (height // 8) * 8  # round down to nearest multiple of 8728 729        if width is None:730            if isinstance(image, PIL.Image.Image):731                width = image.width732            elif isinstance(image, torch.Tensor):733                width = image.shape[2]734 735            width = (width // 8) * 8  # round down to nearest multiple of 8736 737        return height, width738 739    @torch.no_grad()740    @replace_example_docstring(EXAMPLE_DOC_STRING)741    def __call__(742        self,743        prompt: Union[str, List[str]] = None,744        image: Union[torch.Tensor, PIL.Image.Image] = None,745        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,746        controlnet_conditioning_image: Union[747            torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]748        ] = None,749        height: Optional[int] = None,750        width: Optional[int] = None,751        num_inference_steps: int = 50,752        guidance_scale: float = 7.5,753        negative_prompt: Optional[Union[str, List[str]]] = None,754        num_images_per_prompt: Optional[int] = 1,755        eta: float = 0.0,756        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,757        latents: Optional[torch.Tensor] = None,758        prompt_embeds: Optional[torch.Tensor] = None,759        negative_prompt_embeds: Optional[torch.Tensor] = None,760        output_type: Optional[str] = "pil",761        return_dict: bool = True,762        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,763        callback_steps: int = 1,764        cross_attention_kwargs: Optional[Dict[str, Any]] = None,765        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,766    ):767        r"""768        Function invoked when calling the pipeline for generation.769 770        Args:771            prompt (`str` or `List[str]`, *optional*):772                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.773                instead.774            image (`torch.Tensor` or `PIL.Image.Image`):775                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will776                be masked out with `mask_image` and repainted according to `prompt`.777            mask_image (`torch.Tensor` or `PIL.Image.Image`):778                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be779                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted780                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)781                instead of 3, so the expected shape would be `(B, H, W, 1)`.782            controlnet_conditioning_image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]` or `List[PIL.Image.Image]`):783                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If784                the type is specified as `torch.Tensor`, it is passed to ControlNet as is. PIL.Image.Image` can785                also be accepted as an image. The control image is automatically resized to fit the output image.786            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):787                The height in pixels of the generated image.788            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):789                The width in pixels of the generated image.790            num_inference_steps (`int`, *optional*, defaults to 50):791                The number of denoising steps. More denoising steps usually lead to a higher quality image at the792                expense of slower inference.793            guidance_scale (`float`, *optional*, defaults to 7.5):794                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).795                `guidance_scale` is defined as `w` of equation 2. of [Imagen796                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >797                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,798                usually at the expense of lower image quality.799            negative_prompt (`str` or `List[str]`, *optional*):800                The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead.801                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).802            num_images_per_prompt (`int`, *optional*, defaults to 1):803                The number of images to generate per prompt.804            eta (`float`, *optional*, defaults to 0.0):805                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to806                [`schedulers.DDIMScheduler`], will be ignored for others.807            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):808                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)809                to make generation deterministic.810            latents (`torch.Tensor`, *optional*):811                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image812                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents813                tensor will ge generated by sampling using the supplied random `generator`.814            prompt_embeds (`torch.Tensor`, *optional*):815                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not816                provided, text embeddings will be generated from `prompt` input argument.817            negative_prompt_embeds (`torch.Tensor`, *optional*):818                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt819                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input820                argument.821            output_type (`str`, *optional*, defaults to `"pil"`):822                The output format of the generate image. Choose between823                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.824            return_dict (`bool`, *optional*, defaults to `True`):825                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a826                plain tuple.827            callback (`Callable`, *optional*):828                A function that will be called every `callback_steps` steps during inference. The function will be829                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.830            callback_steps (`int`, *optional*, defaults to 1):831                The frequency at which the `callback` function will be called. If not specified, the callback will be832                called at every step.833            cross_attention_kwargs (`dict`, *optional*):834                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under835                `self.processor` in836                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).837            controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):838                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added839                to the residual in the original unet.840 841        Examples:842 843        Returns:844            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:845            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.846            When returning a tuple, the first element is a list with the generated images, and the second element is a847            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"848            (nsfw) content, according to the `safety_checker`.849        """850        # 0. Default height and width to unet851        height, width = self._default_height_width(height, width, controlnet_conditioning_image)852 853        # 1. Check inputs. Raise error if not correct854        self.check_inputs(855            prompt,856            image,857            mask_image,858            controlnet_conditioning_image,859            height,860            width,861            callback_steps,862            negative_prompt,863            prompt_embeds,864            negative_prompt_embeds,865            controlnet_conditioning_scale,866        )867 868        # 2. Define call parameters869        if prompt is not None and isinstance(prompt, str):870            batch_size = 1871        elif prompt is not None and isinstance(prompt, list):872            batch_size = len(prompt)873        else:874            batch_size = prompt_embeds.shape[0]875 876        device = self._execution_device877        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)878        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`879        # corresponds to doing no classifier free guidance.880        do_classifier_free_guidance = guidance_scale > 1.0881 882        if isinstance(self.controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):883            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(self.controlnet.nets)884 885        # 3. Encode input prompt886        prompt_embeds = self._encode_prompt(887            prompt,888            device,889            num_images_per_prompt,890            do_classifier_free_guidance,891            negative_prompt,892            prompt_embeds=prompt_embeds,893            negative_prompt_embeds=negative_prompt_embeds,894        )895 896        # 4. Prepare mask, image, and controlnet_conditioning_image897        image = prepare_image(image)898 899        mask_image = prepare_mask_image(mask_image)900 901        # condition image(s)902        if isinstance(self.controlnet, ControlNetModel):903            controlnet_conditioning_image = prepare_controlnet_conditioning_image(904                controlnet_conditioning_image=controlnet_conditioning_image,905                width=width,906                height=height,907                batch_size=batch_size * num_images_per_prompt,908                num_images_per_prompt=num_images_per_prompt,909                device=device,910                dtype=self.controlnet.dtype,911                do_classifier_free_guidance=do_classifier_free_guidance,912            )913        elif isinstance(self.controlnet, MultiControlNetModel):914            controlnet_conditioning_images = []915 916            for image_ in controlnet_conditioning_image:917                image_ = prepare_controlnet_conditioning_image(918                    controlnet_conditioning_image=image_,919                    width=width,920                    height=height,921                    batch_size=batch_size * num_images_per_prompt,922                    num_images_per_prompt=num_images_per_prompt,923                    device=device,924                    dtype=self.controlnet.dtype,925                    do_classifier_free_guidance=do_classifier_free_guidance,926                )927                controlnet_conditioning_images.append(image_)928 929            controlnet_conditioning_image = controlnet_conditioning_images930        else:931            assert False932 933        masked_image = image * (mask_image < 0.5)934 935        # 5. Prepare timesteps936        self.scheduler.set_timesteps(num_inference_steps, device=device)937        timesteps = self.scheduler.timesteps938 939        # 6. Prepare latent variables940        num_channels_latents = self.vae.config.latent_channels941        latents = self.prepare_latents(942            batch_size * num_images_per_prompt,943            num_channels_latents,944            height,945            width,946            prompt_embeds.dtype,947            device,948            generator,949            latents,950        )951 952        mask_image_latents = self.prepare_mask_latents(953            mask_image,954            batch_size * num_images_per_prompt,955            height,956            width,957            prompt_embeds.dtype,958            device,959            do_classifier_free_guidance,960        )961 962        masked_image_latents = self.prepare_masked_image_latents(963            masked_image,964            batch_size * num_images_per_prompt,965            height,966            width,967            prompt_embeds.dtype,968            device,969            generator,970            do_classifier_free_guidance,971        )972 973        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline974        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)975 976        # 8. Denoising loop977        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order978        with self.progress_bar(total=num_inference_steps) as progress_bar:979            for i, t in enumerate(timesteps):980                # expand the latents if we are doing classifier free guidance981                non_inpainting_latent_model_input = (982                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents983                )984 985                non_inpainting_latent_model_input = self.scheduler.scale_model_input(986                    non_inpainting_latent_model_input, t987                )988 989                inpainting_latent_model_input = torch.cat(990                    [non_inpainting_latent_model_input, mask_image_latents, masked_image_latents], dim=1991                )992 993                down_block_res_samples, mid_block_res_sample = self.controlnet(994                    non_inpainting_latent_model_input,995                    t,996                    encoder_hidden_states=prompt_embeds,997                    controlnet_cond=controlnet_conditioning_image,998                    conditioning_scale=controlnet_conditioning_scale,999                    return_dict=False,1000                )1001 1002                # predict the noise residual1003                noise_pred = self.unet(1004                    inpainting_latent_model_input,1005                    t,1006                    encoder_hidden_states=prompt_embeds,1007                    cross_attention_kwargs=cross_attention_kwargs,1008                    down_block_additional_residuals=down_block_res_samples,1009                    mid_block_additional_residual=mid_block_res_sample,1010                ).sample1011 1012                # perform guidance1013                if do_classifier_free_guidance:1014                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)1015                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)1016 1017                # compute the previous noisy sample x_t -> x_t-11018                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample1019 1020                # call the callback, if provided1021                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):1022                    progress_bar.update()1023                    if callback is not None and i % callback_steps == 0:1024                        step_idx = i // getattr(self.scheduler, "order", 1)1025                        callback(step_idx, t, latents)1026 1027        # If we do sequential model offloading, let's offload unet and controlnet1028        # manually for max memory savings1029        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:1030            self.unet.to("cpu")1031            self.controlnet.to("cpu")1032            torch.cuda.empty_cache()1033 1034        if output_type == "latent":1035            image = latents1036            has_nsfw_concept = None1037        elif output_type == "pil":1038            # 8. Post-processing1039            image = self.decode_latents(latents)1040 1041            # 9. Run safety checker1042            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)1043 1044            # 10. Convert to PIL1045            image = self.numpy_to_pil(image)1046        else:1047            # 8. Post-processing1048            image = self.decode_latents(latents)1049 1050            # 9. Run safety checker1051            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)1052 1053        # Offload last model to CPU1054        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:1055            self.final_offload_hook.offload()1056 1057        if not return_dict:1058            return (image, has_nsfw_concept)1059 1060        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)1061