diffusers/community-pipelines-mirror

1# Inspired by: https://github.com/Mikubill/sd-webui-controlnet/discussions/1236 and https://github.com/Mikubill/sd-webui-controlnet/discussions/12802from typing import Any, Callable, Dict, List, Optional, Tuple, Union3 4import numpy as np5import PIL.Image6import torch7 8from diffusers import StableDiffusionControlNetPipeline9from diffusers.models import ControlNetModel10from diffusers.models.attention import BasicTransformerBlock11from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D12from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel13from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput14from diffusers.utils import is_compiled_module, logging, randn_tensor15 16 17logger = logging.get_logger(__name__)  # pylint: disable=invalid-name18 19EXAMPLE_DOC_STRING = """20    Examples:21        ```py22        >>> import cv223        >>> import torch24        >>> import numpy as np25        >>> from PIL import Image26        >>> from diffusers import UniPCMultistepScheduler27        >>> from diffusers.utils import load_image28 29        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")30 31        >>> # get canny image32        >>> image = cv2.Canny(np.array(input_image), 100, 200)33        >>> image = image[:, :, None]34        >>> image = np.concatenate([image, image, image], axis=2)35        >>> canny_image = Image.fromarray(image)36 37        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)38        >>> pipe = StableDiffusionControlNetReferencePipeline.from_pretrained(39                "runwayml/stable-diffusion-v1-5",40                controlnet=controlnet,41                safety_checker=None,42                torch_dtype=torch.float1643                ).to('cuda:0')44 45        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)46 47        >>> result_img = pipe(ref_image=input_image,48                        prompt="1girl",49                        image=canny_image,50                        num_inference_steps=20,51                        reference_attn=True,52                        reference_adain=True).images[0]53 54        >>> result_img.show()55        ```56"""57 58 59def torch_dfs(model: torch.nn.Module):60    result = [model]61    for child in model.children():62        result += torch_dfs(child)63    return result64 65 66class StableDiffusionControlNetReferencePipeline(StableDiffusionControlNetPipeline):67    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):68        refimage = refimage.to(device=device, dtype=dtype)69 70        # encode the mask image into latents space so we can concatenate it to the latents71        if isinstance(generator, list):72            ref_image_latents = [73                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])74                for i in range(batch_size)75            ]76            ref_image_latents = torch.cat(ref_image_latents, dim=0)77        else:78            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)79        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents80 81        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method82        if ref_image_latents.shape[0] < batch_size:83            if not batch_size % ref_image_latents.shape[0] == 0:84                raise ValueError(85                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"86                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."87                    " Make sure the number of images that you pass is divisible by the total requested batch size."88                )89            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)90 91        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents92 93        # aligning device to prevent device errors when concating it with the latent model input94        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)95        return ref_image_latents96 97    @torch.no_grad()98    def __call__(99        self,100        prompt: Union[str, List[str]] = None,101        image: Union[102            torch.FloatTensor,103            PIL.Image.Image,104            np.ndarray,105            List[torch.FloatTensor],106            List[PIL.Image.Image],107            List[np.ndarray],108        ] = None,109        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,110        height: Optional[int] = None,111        width: Optional[int] = None,112        num_inference_steps: int = 50,113        guidance_scale: float = 7.5,114        negative_prompt: Optional[Union[str, List[str]]] = None,115        num_images_per_prompt: Optional[int] = 1,116        eta: float = 0.0,117        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,118        latents: Optional[torch.FloatTensor] = None,119        prompt_embeds: Optional[torch.FloatTensor] = None,120        negative_prompt_embeds: Optional[torch.FloatTensor] = None,121        output_type: Optional[str] = "pil",122        return_dict: bool = True,123        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,124        callback_steps: int = 1,125        cross_attention_kwargs: Optional[Dict[str, Any]] = None,126        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,127        guess_mode: bool = False,128        attention_auto_machine_weight: float = 1.0,129        gn_auto_machine_weight: float = 1.0,130        style_fidelity: float = 0.5,131        reference_attn: bool = True,132        reference_adain: bool = True,133    ):134        r"""135        Function invoked when calling the pipeline for generation.136 137        Args:138            prompt (`str` or `List[str]`, *optional*):139                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.140                instead.141            image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:142                    `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):143                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If144                the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can145                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If146                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are147                specified in init, images must be passed as a list such that each element of the list can be correctly148                batched for input to a single controlnet.149            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):150                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If151                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can152                also be accepted as an image.153            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):154                The height in pixels of the generated image.155            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):156                The width in pixels of the generated image.157            num_inference_steps (`int`, *optional*, defaults to 50):158                The number of denoising steps. More denoising steps usually lead to a higher quality image at the159                expense of slower inference.160            guidance_scale (`float`, *optional*, defaults to 7.5):161                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).162                `guidance_scale` is defined as `w` of equation 2. of [Imagen163                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >164                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,165                usually at the expense of lower image quality.166            negative_prompt (`str` or `List[str]`, *optional*):167                The prompt or prompts not to guide the image generation. If not defined, one has to pass168                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is169                less than `1`).170            num_images_per_prompt (`int`, *optional*, defaults to 1):171                The number of images to generate per prompt.172            eta (`float`, *optional*, defaults to 0.0):173                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to174                [`schedulers.DDIMScheduler`], will be ignored for others.175            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):176                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)177                to make generation deterministic.178            latents (`torch.FloatTensor`, *optional*):179                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image180                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents181                tensor will ge generated by sampling using the supplied random `generator`.182            prompt_embeds (`torch.FloatTensor`, *optional*):183                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not184                provided, text embeddings will be generated from `prompt` input argument.185            negative_prompt_embeds (`torch.FloatTensor`, *optional*):186                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt187                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input188                argument.189            output_type (`str`, *optional*, defaults to `"pil"`):190                The output format of the generate image. Choose between191                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.192            return_dict (`bool`, *optional*, defaults to `True`):193                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a194                plain tuple.195            callback (`Callable`, *optional*):196                A function that will be called every `callback_steps` steps during inference. The function will be197                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.198            callback_steps (`int`, *optional*, defaults to 1):199                The frequency at which the `callback` function will be called. If not specified, the callback will be200                called at every step.201            cross_attention_kwargs (`dict`, *optional*):202                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under203                `self.processor` in204                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).205            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):206                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added207                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the208                corresponding scale as a list.209            guess_mode (`bool`, *optional*, defaults to `False`):210                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if211                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.212            attention_auto_machine_weight (`float`):213                Weight of using reference query for self attention's context.214                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.215            gn_auto_machine_weight (`float`):216                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.217            style_fidelity (`float`):218                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,219                elif style_fidelity=0.0, prompt more important, else balanced.220            reference_attn (`bool`):221                Whether to use reference query for self attention's context.222            reference_adain (`bool`):223                Whether to use reference adain.224 225        Examples:226 227        Returns:228            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:229            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.230            When returning a tuple, the first element is a list with the generated images, and the second element is a231            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"232            (nsfw) content, according to the `safety_checker`.233        """234        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."235 236        # 1. Check inputs. Raise error if not correct237        self.check_inputs(238            prompt,239            image,240            callback_steps,241            negative_prompt,242            prompt_embeds,243            negative_prompt_embeds,244            controlnet_conditioning_scale,245        )246 247        # 2. Define call parameters248        if prompt is not None and isinstance(prompt, str):249            batch_size = 1250        elif prompt is not None and isinstance(prompt, list):251            batch_size = len(prompt)252        else:253            batch_size = prompt_embeds.shape[0]254 255        device = self._execution_device256        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)257        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`258        # corresponds to doing no classifier free guidance.259        do_classifier_free_guidance = guidance_scale > 1.0260 261        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet262 263        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):264            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)265 266        global_pool_conditions = (267            controlnet.config.global_pool_conditions268            if isinstance(controlnet, ControlNetModel)269            else controlnet.nets[0].config.global_pool_conditions270        )271        guess_mode = guess_mode or global_pool_conditions272 273        # 3. Encode input prompt274        text_encoder_lora_scale = (275            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None276        )277        prompt_embeds = self._encode_prompt(278            prompt,279            device,280            num_images_per_prompt,281            do_classifier_free_guidance,282            negative_prompt,283            prompt_embeds=prompt_embeds,284            negative_prompt_embeds=negative_prompt_embeds,285            lora_scale=text_encoder_lora_scale,286        )287 288        # 4. Prepare image289        if isinstance(controlnet, ControlNetModel):290            image = self.prepare_image(291                image=image,292                width=width,293                height=height,294                batch_size=batch_size * num_images_per_prompt,295                num_images_per_prompt=num_images_per_prompt,296                device=device,297                dtype=controlnet.dtype,298                do_classifier_free_guidance=do_classifier_free_guidance,299                guess_mode=guess_mode,300            )301            height, width = image.shape[-2:]302        elif isinstance(controlnet, MultiControlNetModel):303            images = []304 305            for image_ in image:306                image_ = self.prepare_image(307                    image=image_,308                    width=width,309                    height=height,310                    batch_size=batch_size * num_images_per_prompt,311                    num_images_per_prompt=num_images_per_prompt,312                    device=device,313                    dtype=controlnet.dtype,314                    do_classifier_free_guidance=do_classifier_free_guidance,315                    guess_mode=guess_mode,316                )317 318                images.append(image_)319 320            image = images321            height, width = image[0].shape[-2:]322        else:323            assert False324 325        # 5. Preprocess reference image326        ref_image = self.prepare_image(327            image=ref_image,328            width=width,329            height=height,330            batch_size=batch_size * num_images_per_prompt,331            num_images_per_prompt=num_images_per_prompt,332            device=device,333            dtype=prompt_embeds.dtype,334        )335 336        # 6. Prepare timesteps337        self.scheduler.set_timesteps(num_inference_steps, device=device)338        timesteps = self.scheduler.timesteps339 340        # 7. Prepare latent variables341        num_channels_latents = self.unet.config.in_channels342        latents = self.prepare_latents(343            batch_size * num_images_per_prompt,344            num_channels_latents,345            height,346            width,347            prompt_embeds.dtype,348            device,349            generator,350            latents,351        )352 353        # 8. Prepare reference latent variables354        ref_image_latents = self.prepare_ref_latents(355            ref_image,356            batch_size * num_images_per_prompt,357            prompt_embeds.dtype,358            device,359            generator,360            do_classifier_free_guidance,361        )362 363        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline364        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)365 366        # 9. Modify self attention and group norm367        MODE = "write"368        uc_mask = (369            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)370            .type_as(ref_image_latents)371            .bool()372        )373 374        def hacked_basic_transformer_inner_forward(375            self,376            hidden_states: torch.FloatTensor,377            attention_mask: Optional[torch.FloatTensor] = None,378            encoder_hidden_states: Optional[torch.FloatTensor] = None,379            encoder_attention_mask: Optional[torch.FloatTensor] = None,380            timestep: Optional[torch.LongTensor] = None,381            cross_attention_kwargs: Dict[str, Any] = None,382            class_labels: Optional[torch.LongTensor] = None,383        ):384            if self.use_ada_layer_norm:385                norm_hidden_states = self.norm1(hidden_states, timestep)386            elif self.use_ada_layer_norm_zero:387                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(388                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype389                )390            else:391                norm_hidden_states = self.norm1(hidden_states)392 393            # 1. Self-Attention394            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}395            if self.only_cross_attention:396                attn_output = self.attn1(397                    norm_hidden_states,398                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,399                    attention_mask=attention_mask,400                    **cross_attention_kwargs,401                )402            else:403                if MODE == "write":404                    self.bank.append(norm_hidden_states.detach().clone())405                    attn_output = self.attn1(406                        norm_hidden_states,407                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,408                        attention_mask=attention_mask,409                        **cross_attention_kwargs,410                    )411                if MODE == "read":412                    if attention_auto_machine_weight > self.attn_weight:413                        attn_output_uc = self.attn1(414                            norm_hidden_states,415                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),416                            # attention_mask=attention_mask,417                            **cross_attention_kwargs,418                        )419                        attn_output_c = attn_output_uc.clone()420                        if do_classifier_free_guidance and style_fidelity > 0:421                            attn_output_c[uc_mask] = self.attn1(422                                norm_hidden_states[uc_mask],423                                encoder_hidden_states=norm_hidden_states[uc_mask],424                                **cross_attention_kwargs,425                            )426                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc427                        self.bank.clear()428                    else:429                        attn_output = self.attn1(430                            norm_hidden_states,431                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,432                            attention_mask=attention_mask,433                            **cross_attention_kwargs,434                        )435            if self.use_ada_layer_norm_zero:436                attn_output = gate_msa.unsqueeze(1) * attn_output437            hidden_states = attn_output + hidden_states438 439            if self.attn2 is not None:440                norm_hidden_states = (441                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)442                )443 444                # 2. Cross-Attention445                attn_output = self.attn2(446                    norm_hidden_states,447                    encoder_hidden_states=encoder_hidden_states,448                    attention_mask=encoder_attention_mask,449                    **cross_attention_kwargs,450                )451                hidden_states = attn_output + hidden_states452 453            # 3. Feed-forward454            norm_hidden_states = self.norm3(hidden_states)455 456            if self.use_ada_layer_norm_zero:457                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]458 459            ff_output = self.ff(norm_hidden_states)460 461            if self.use_ada_layer_norm_zero:462                ff_output = gate_mlp.unsqueeze(1) * ff_output463 464            hidden_states = ff_output + hidden_states465 466            return hidden_states467 468        def hacked_mid_forward(self, *args, **kwargs):469            eps = 1e-6470            x = self.original_forward(*args, **kwargs)471            if MODE == "write":472                if gn_auto_machine_weight >= self.gn_weight:473                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)474                    self.mean_bank.append(mean)475                    self.var_bank.append(var)476            if MODE == "read":477                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:478                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)479                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5480                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))481                    var_acc = sum(self.var_bank) / float(len(self.var_bank))482                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5483                    x_uc = (((x - mean) / std) * std_acc) + mean_acc484                    x_c = x_uc.clone()485                    if do_classifier_free_guidance and style_fidelity > 0:486                        x_c[uc_mask] = x[uc_mask]487                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc488                self.mean_bank = []489                self.var_bank = []490            return x491 492        def hack_CrossAttnDownBlock2D_forward(493            self,494            hidden_states: torch.FloatTensor,495            temb: Optional[torch.FloatTensor] = None,496            encoder_hidden_states: Optional[torch.FloatTensor] = None,497            attention_mask: Optional[torch.FloatTensor] = None,498            cross_attention_kwargs: Optional[Dict[str, Any]] = None,499            encoder_attention_mask: Optional[torch.FloatTensor] = None,500        ):501            eps = 1e-6502 503            # TODO(Patrick, William) - attention mask is not used504            output_states = ()505 506            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):507                hidden_states = resnet(hidden_states, temb)508                hidden_states = attn(509                    hidden_states,510                    encoder_hidden_states=encoder_hidden_states,511                    cross_attention_kwargs=cross_attention_kwargs,512                    attention_mask=attention_mask,513                    encoder_attention_mask=encoder_attention_mask,514                    return_dict=False,515                )[0]516                if MODE == "write":517                    if gn_auto_machine_weight >= self.gn_weight:518                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)519                        self.mean_bank.append([mean])520                        self.var_bank.append([var])521                if MODE == "read":522                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:523                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)524                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5525                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))526                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))527                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5528                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc529                        hidden_states_c = hidden_states_uc.clone()530                        if do_classifier_free_guidance and style_fidelity > 0:531                            hidden_states_c[uc_mask] = hidden_states[uc_mask]532                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc533 534                output_states = output_states + (hidden_states,)535 536            if MODE == "read":537                self.mean_bank = []538                self.var_bank = []539 540            if self.downsamplers is not None:541                for downsampler in self.downsamplers:542                    hidden_states = downsampler(hidden_states)543 544                output_states = output_states + (hidden_states,)545 546            return hidden_states, output_states547 548        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):549            eps = 1e-6550 551            output_states = ()552 553            for i, resnet in enumerate(self.resnets):554                hidden_states = resnet(hidden_states, temb)555 556                if MODE == "write":557                    if gn_auto_machine_weight >= self.gn_weight:558                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)559                        self.mean_bank.append([mean])560                        self.var_bank.append([var])561                if MODE == "read":562                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:563                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)564                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5565                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))566                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))567                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5568                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc569                        hidden_states_c = hidden_states_uc.clone()570                        if do_classifier_free_guidance and style_fidelity > 0:571                            hidden_states_c[uc_mask] = hidden_states[uc_mask]572                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc573 574                output_states = output_states + (hidden_states,)575 576            if MODE == "read":577                self.mean_bank = []578                self.var_bank = []579 580            if self.downsamplers is not None:581                for downsampler in self.downsamplers:582                    hidden_states = downsampler(hidden_states)583 584                output_states = output_states + (hidden_states,)585 586            return hidden_states, output_states587 588        def hacked_CrossAttnUpBlock2D_forward(589            self,590            hidden_states: torch.FloatTensor,591            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],592            temb: Optional[torch.FloatTensor] = None,593            encoder_hidden_states: Optional[torch.FloatTensor] = None,594            cross_attention_kwargs: Optional[Dict[str, Any]] = None,595            upsample_size: Optional[int] = None,596            attention_mask: Optional[torch.FloatTensor] = None,597            encoder_attention_mask: Optional[torch.FloatTensor] = None,598        ):599            eps = 1e-6600            # TODO(Patrick, William) - attention mask is not used601            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):602                # pop res hidden states603                res_hidden_states = res_hidden_states_tuple[-1]604                res_hidden_states_tuple = res_hidden_states_tuple[:-1]605                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)606                hidden_states = resnet(hidden_states, temb)607                hidden_states = attn(608                    hidden_states,609                    encoder_hidden_states=encoder_hidden_states,610                    cross_attention_kwargs=cross_attention_kwargs,611                    attention_mask=attention_mask,612                    encoder_attention_mask=encoder_attention_mask,613                    return_dict=False,614                )[0]615 616                if MODE == "write":617                    if gn_auto_machine_weight >= self.gn_weight:618                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)619                        self.mean_bank.append([mean])620                        self.var_bank.append([var])621                if MODE == "read":622                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:623                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)624                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5625                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))626                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))627                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5628                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc629                        hidden_states_c = hidden_states_uc.clone()630                        if do_classifier_free_guidance and style_fidelity > 0:631                            hidden_states_c[uc_mask] = hidden_states[uc_mask]632                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc633 634            if MODE == "read":635                self.mean_bank = []636                self.var_bank = []637 638            if self.upsamplers is not None:639                for upsampler in self.upsamplers:640                    hidden_states = upsampler(hidden_states, upsample_size)641 642            return hidden_states643 644        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):645            eps = 1e-6646            for i, resnet in enumerate(self.resnets):647                # pop res hidden states648                res_hidden_states = res_hidden_states_tuple[-1]649                res_hidden_states_tuple = res_hidden_states_tuple[:-1]650                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)651                hidden_states = resnet(hidden_states, temb)652 653                if MODE == "write":654                    if gn_auto_machine_weight >= self.gn_weight:655                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)656                        self.mean_bank.append([mean])657                        self.var_bank.append([var])658                if MODE == "read":659                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:660                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)661                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5662                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))663                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))664                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5665                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc666                        hidden_states_c = hidden_states_uc.clone()667                        if do_classifier_free_guidance and style_fidelity > 0:668                            hidden_states_c[uc_mask] = hidden_states[uc_mask]669                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc670 671            if MODE == "read":672                self.mean_bank = []673                self.var_bank = []674 675            if self.upsamplers is not None:676                for upsampler in self.upsamplers:677                    hidden_states = upsampler(hidden_states, upsample_size)678 679            return hidden_states680 681        if reference_attn:682            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]683            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])684 685            for i, module in enumerate(attn_modules):686                module._original_inner_forward = module.forward687                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)688                module.bank = []689                module.attn_weight = float(i) / float(len(attn_modules))690 691        if reference_adain:692            gn_modules = [self.unet.mid_block]693            self.unet.mid_block.gn_weight = 0694 695            down_blocks = self.unet.down_blocks696            for w, module in enumerate(down_blocks):697                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))698                gn_modules.append(module)699 700            up_blocks = self.unet.up_blocks701            for w, module in enumerate(up_blocks):702                module.gn_weight = float(w) / float(len(up_blocks))703                gn_modules.append(module)704 705            for i, module in enumerate(gn_modules):706                if getattr(module, "original_forward", None) is None:707                    module.original_forward = module.forward708                if i == 0:709                    # mid_block710                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)711                elif isinstance(module, CrossAttnDownBlock2D):712                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)713                elif isinstance(module, DownBlock2D):714                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)715                elif isinstance(module, CrossAttnUpBlock2D):716                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)717                elif isinstance(module, UpBlock2D):718                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)719                module.mean_bank = []720                module.var_bank = []721                module.gn_weight *= 2722 723        # 11. Denoising loop724        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order725        with self.progress_bar(total=num_inference_steps) as progress_bar:726            for i, t in enumerate(timesteps):727                # expand the latents if we are doing classifier free guidance728                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents729                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)730 731                # controlnet(s) inference732                if guess_mode and do_classifier_free_guidance:733                    # Infer ControlNet only for the conditional batch.734                    control_model_input = latents735                    control_model_input = self.scheduler.scale_model_input(control_model_input, t)736                    controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]737                else:738                    control_model_input = latent_model_input739                    controlnet_prompt_embeds = prompt_embeds740 741                down_block_res_samples, mid_block_res_sample = self.controlnet(742                    control_model_input,743                    t,744                    encoder_hidden_states=controlnet_prompt_embeds,745                    controlnet_cond=image,746                    conditioning_scale=controlnet_conditioning_scale,747                    guess_mode=guess_mode,748                    return_dict=False,749                )750 751                if guess_mode and do_classifier_free_guidance:752                    # Infered ControlNet only for the conditional batch.753                    # To apply the output of ControlNet to both the unconditional and conditional batches,754                    # add 0 to the unconditional batch to keep it unchanged.755                    down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]756                    mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])757 758                # ref only part759                noise = randn_tensor(760                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype761                )762                ref_xt = self.scheduler.add_noise(763                    ref_image_latents,764                    noise,765                    t.reshape(766                        1,767                    ),768                )769                ref_xt = self.scheduler.scale_model_input(ref_xt, t)770 771                MODE = "write"772                self.unet(773                    ref_xt,774                    t,775                    encoder_hidden_states=prompt_embeds,776                    cross_attention_kwargs=cross_attention_kwargs,777                    return_dict=False,778                )779 780                # predict the noise residual781                MODE = "read"782                noise_pred = self.unet(783                    latent_model_input,784                    t,785                    encoder_hidden_states=prompt_embeds,786                    cross_attention_kwargs=cross_attention_kwargs,787                    down_block_additional_residuals=down_block_res_samples,788                    mid_block_additional_residual=mid_block_res_sample,789                    return_dict=False,790                )[0]791 792                # perform guidance793                if do_classifier_free_guidance:794                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)795                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)796 797                # compute the previous noisy sample x_t -> x_t-1798                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]799 800                # call the callback, if provided801                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):802                    progress_bar.update()803                    if callback is not None and i % callback_steps == 0:804                        callback(i, t, latents)805 806        # If we do sequential model offloading, let's offload unet and controlnet807        # manually for max memory savings808        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:809            self.unet.to("cpu")810            self.controlnet.to("cpu")811            torch.cuda.empty_cache()812 813        if not output_type == "latent":814            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]815            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)816        else:817            image = latents818            has_nsfw_concept = None819 820        if has_nsfw_concept is None:821            do_denormalize = [True] * image.shape[0]822        else:823            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]824 825        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)826 827        # Offload last model to CPU828        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:829            self.final_offload_hook.offload()830 831        if not return_dict:832            return (image, has_nsfw_concept)833 834        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)835