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 PIL.Image5import torch6 7from diffusers import StableDiffusionControlNetPipeline8from diffusers.models import ControlNetModel9from diffusers.models.attention import BasicTransformerBlock10from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D11from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel12from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput13from diffusers.utils import is_compiled_module, logging, randn_tensor14 15 16logger = logging.get_logger(__name__)  # pylint: disable=invalid-name17 18EXAMPLE_DOC_STRING = """19    Examples:20        ```py21        >>> import cv222        >>> import torch23        >>> import numpy as np24        >>> from PIL import Image25        >>> from diffusers import UniPCMultistepScheduler26        >>> from diffusers.utils import load_image27 28        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")29 30        >>> # get canny image31        >>> image = cv2.Canny(np.array(input_image), 100, 200)32        >>> image = image[:, :, None]33        >>> image = np.concatenate([image, image, image], axis=2)34        >>> canny_image = Image.fromarray(image)35 36        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)37        >>> pipe = StableDiffusionControlNetReferencePipeline.from_pretrained(38                "runwayml/stable-diffusion-v1-5",39                controlnet=controlnet,40                safety_checker=None,41                torch_dtype=torch.float1642                ).to('cuda:0')43 44        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)45 46        >>> result_img = pipe(ref_image=input_image,47                        prompt="1girl",48                        image=canny_image,49                        num_inference_steps=20,50                        reference_attn=True,51                        reference_adain=True).images[0]52 53        >>> result_img.show()54        ```55"""56 57 58def torch_dfs(model: torch.nn.Module):59    result = [model]60    for child in model.children():61        result += torch_dfs(child)62    return result63 64 65class StableDiffusionControlNetReferencePipeline(StableDiffusionControlNetPipeline):66    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):67        refimage = refimage.to(device=device, dtype=dtype)68 69        # encode the mask image into latents space so we can concatenate it to the latents70        if isinstance(generator, list):71            ref_image_latents = [72                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])73                for i in range(batch_size)74            ]75            ref_image_latents = torch.cat(ref_image_latents, dim=0)76        else:77            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)78        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents79 80        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method81        if ref_image_latents.shape[0] < batch_size:82            if not batch_size % ref_image_latents.shape[0] == 0:83                raise ValueError(84                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"85                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."86                    " Make sure the number of images that you pass is divisible by the total requested batch size."87                )88            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)89 90        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents91 92        # aligning device to prevent device errors when concating it with the latent model input93        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)94        return ref_image_latents95 96    @torch.no_grad()97    def __call__(98        self,99        prompt: Union[str, List[str]] = None,100        image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None,101        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,102        height: Optional[int] = None,103        width: Optional[int] = None,104        num_inference_steps: int = 50,105        guidance_scale: float = 7.5,106        negative_prompt: Optional[Union[str, List[str]]] = None,107        num_images_per_prompt: Optional[int] = 1,108        eta: float = 0.0,109        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,110        latents: Optional[torch.FloatTensor] = None,111        prompt_embeds: Optional[torch.FloatTensor] = None,112        negative_prompt_embeds: Optional[torch.FloatTensor] = None,113        output_type: Optional[str] = "pil",114        return_dict: bool = True,115        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,116        callback_steps: int = 1,117        cross_attention_kwargs: Optional[Dict[str, Any]] = None,118        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,119        guess_mode: bool = False,120        attention_auto_machine_weight: float = 1.0,121        gn_auto_machine_weight: float = 1.0,122        style_fidelity: float = 0.5,123        reference_attn: bool = True,124        reference_adain: bool = True,125    ):126        r"""127        Function invoked when calling the pipeline for generation.128 129        Args:130            prompt (`str` or `List[str]`, *optional*):131                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.132                instead.133            image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`,134                    `List[List[torch.FloatTensor]]`, or `List[List[PIL.Image.Image]]`):135                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If136                the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can137                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If138                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are139                specified in init, images must be passed as a list such that each element of the list can be correctly140                batched for input to a single controlnet.141            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):142                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If143                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can144                also be accepted as an image.145            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):146                The height in pixels of the generated image.147            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):148                The width in pixels of the generated image.149            num_inference_steps (`int`, *optional*, defaults to 50):150                The number of denoising steps. More denoising steps usually lead to a higher quality image at the151                expense of slower inference.152            guidance_scale (`float`, *optional*, defaults to 7.5):153                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).154                `guidance_scale` is defined as `w` of equation 2. of [Imagen155                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >156                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,157                usually at the expense of lower image quality.158            negative_prompt (`str` or `List[str]`, *optional*):159                The prompt or prompts not to guide the image generation. If not defined, one has to pass160                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is161                less than `1`).162            num_images_per_prompt (`int`, *optional*, defaults to 1):163                The number of images to generate per prompt.164            eta (`float`, *optional*, defaults to 0.0):165                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to166                [`schedulers.DDIMScheduler`], will be ignored for others.167            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):168                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)169                to make generation deterministic.170            latents (`torch.FloatTensor`, *optional*):171                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image172                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents173                tensor will ge generated by sampling using the supplied random `generator`.174            prompt_embeds (`torch.FloatTensor`, *optional*):175                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not176                provided, text embeddings will be generated from `prompt` input argument.177            negative_prompt_embeds (`torch.FloatTensor`, *optional*):178                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt179                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input180                argument.181            output_type (`str`, *optional*, defaults to `"pil"`):182                The output format of the generate image. Choose between183                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.184            return_dict (`bool`, *optional*, defaults to `True`):185                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a186                plain tuple.187            callback (`Callable`, *optional*):188                A function that will be called every `callback_steps` steps during inference. The function will be189                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.190            callback_steps (`int`, *optional*, defaults to 1):191                The frequency at which the `callback` function will be called. If not specified, the callback will be192                called at every step.193            cross_attention_kwargs (`dict`, *optional*):194                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under195                `self.processor` in196                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).197            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):198                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added199                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the200                corresponding scale as a list.201            guess_mode (`bool`, *optional*, defaults to `False`):202                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if203                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.204            attention_auto_machine_weight (`float`):205                Weight of using reference query for self attention's context.206                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.207            gn_auto_machine_weight (`float`):208                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.209            style_fidelity (`float`):210                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,211                elif style_fidelity=0.0, prompt more important, else balanced.212            reference_attn (`bool`):213                Whether to use reference query for self attention's context.214            reference_adain (`bool`):215                Whether to use reference adain.216 217        Examples:218 219        Returns:220            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:221            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.222            When returning a tuple, the first element is a list with the generated images, and the second element is a223            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"224            (nsfw) content, according to the `safety_checker`.225        """226        # 0. Default height and width to unet227        height, width = self._default_height_width(height, width, image)228 229        # 1. Check inputs. Raise error if not correct230        self.check_inputs(231            prompt,232            image,233            height,234            width,235            callback_steps,236            negative_prompt,237            prompt_embeds,238            negative_prompt_embeds,239            controlnet_conditioning_scale,240        )241 242        # 2. Define call parameters243        if prompt is not None and isinstance(prompt, str):244            batch_size = 1245        elif prompt is not None and isinstance(prompt, list):246            batch_size = len(prompt)247        else:248            batch_size = prompt_embeds.shape[0]249 250        device = self._execution_device251        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)252        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`253        # corresponds to doing no classifier free guidance.254        do_classifier_free_guidance = guidance_scale > 1.0255 256        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet257 258        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):259            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)260 261        global_pool_conditions = (262            controlnet.config.global_pool_conditions263            if isinstance(controlnet, ControlNetModel)264            else controlnet.nets[0].config.global_pool_conditions265        )266        guess_mode = guess_mode or global_pool_conditions267 268        # 3. Encode input prompt269        prompt_embeds = self._encode_prompt(270            prompt,271            device,272            num_images_per_prompt,273            do_classifier_free_guidance,274            negative_prompt,275            prompt_embeds=prompt_embeds,276            negative_prompt_embeds=negative_prompt_embeds,277        )278 279        # 4. Prepare image280        if isinstance(controlnet, ControlNetModel):281            image = self.prepare_image(282                image=image,283                width=width,284                height=height,285                batch_size=batch_size * num_images_per_prompt,286                num_images_per_prompt=num_images_per_prompt,287                device=device,288                dtype=controlnet.dtype,289                do_classifier_free_guidance=do_classifier_free_guidance,290                guess_mode=guess_mode,291            )292        elif isinstance(controlnet, MultiControlNetModel):293            images = []294 295            for image_ in image:296                image_ = self.prepare_image(297                    image=image_,298                    width=width,299                    height=height,300                    batch_size=batch_size * num_images_per_prompt,301                    num_images_per_prompt=num_images_per_prompt,302                    device=device,303                    dtype=controlnet.dtype,304                    do_classifier_free_guidance=do_classifier_free_guidance,305                    guess_mode=guess_mode,306                )307 308                images.append(image_)309 310            image = images311        else:312            assert False313 314        # 5. Preprocess reference image315        ref_image = self.prepare_image(316            image=ref_image,317            width=width,318            height=height,319            batch_size=batch_size * num_images_per_prompt,320            num_images_per_prompt=num_images_per_prompt,321            device=device,322            dtype=prompt_embeds.dtype,323        )324 325        # 6. Prepare timesteps326        self.scheduler.set_timesteps(num_inference_steps, device=device)327        timesteps = self.scheduler.timesteps328 329        # 7. Prepare latent variables330        num_channels_latents = self.unet.config.in_channels331        latents = self.prepare_latents(332            batch_size * num_images_per_prompt,333            num_channels_latents,334            height,335            width,336            prompt_embeds.dtype,337            device,338            generator,339            latents,340        )341 342        # 8. Prepare reference latent variables343        ref_image_latents = self.prepare_ref_latents(344            ref_image,345            batch_size * num_images_per_prompt,346            prompt_embeds.dtype,347            device,348            generator,349            do_classifier_free_guidance,350        )351 352        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline353        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)354 355        # 9. Modify self attention and group norm356        MODE = "write"357        uc_mask = (358            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)359            .type_as(ref_image_latents)360            .bool()361        )362 363        def hacked_basic_transformer_inner_forward(364            self,365            hidden_states: torch.FloatTensor,366            attention_mask: Optional[torch.FloatTensor] = None,367            encoder_hidden_states: Optional[torch.FloatTensor] = None,368            encoder_attention_mask: Optional[torch.FloatTensor] = None,369            timestep: Optional[torch.LongTensor] = None,370            cross_attention_kwargs: Dict[str, Any] = None,371            class_labels: Optional[torch.LongTensor] = None,372        ):373            if self.use_ada_layer_norm:374                norm_hidden_states = self.norm1(hidden_states, timestep)375            elif self.use_ada_layer_norm_zero:376                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(377                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype378                )379            else:380                norm_hidden_states = self.norm1(hidden_states)381 382            # 1. Self-Attention383            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}384            if self.only_cross_attention:385                attn_output = self.attn1(386                    norm_hidden_states,387                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,388                    attention_mask=attention_mask,389                    **cross_attention_kwargs,390                )391            else:392                if MODE == "write":393                    self.bank.append(norm_hidden_states.detach().clone())394                    attn_output = self.attn1(395                        norm_hidden_states,396                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,397                        attention_mask=attention_mask,398                        **cross_attention_kwargs,399                    )400                if MODE == "read":401                    if attention_auto_machine_weight > self.attn_weight:402                        attn_output_uc = self.attn1(403                            norm_hidden_states,404                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),405                            # attention_mask=attention_mask,406                            **cross_attention_kwargs,407                        )408                        attn_output_c = attn_output_uc.clone()409                        if do_classifier_free_guidance and style_fidelity > 0:410                            attn_output_c[uc_mask] = self.attn1(411                                norm_hidden_states[uc_mask],412                                encoder_hidden_states=norm_hidden_states[uc_mask],413                                **cross_attention_kwargs,414                            )415                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc416                        self.bank.clear()417                    else:418                        attn_output = self.attn1(419                            norm_hidden_states,420                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,421                            attention_mask=attention_mask,422                            **cross_attention_kwargs,423                        )424            if self.use_ada_layer_norm_zero:425                attn_output = gate_msa.unsqueeze(1) * attn_output426            hidden_states = attn_output + hidden_states427 428            if self.attn2 is not None:429                norm_hidden_states = (430                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)431                )432 433                # 2. Cross-Attention434                attn_output = self.attn2(435                    norm_hidden_states,436                    encoder_hidden_states=encoder_hidden_states,437                    attention_mask=encoder_attention_mask,438                    **cross_attention_kwargs,439                )440                hidden_states = attn_output + hidden_states441 442            # 3. Feed-forward443            norm_hidden_states = self.norm3(hidden_states)444 445            if self.use_ada_layer_norm_zero:446                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]447 448            ff_output = self.ff(norm_hidden_states)449 450            if self.use_ada_layer_norm_zero:451                ff_output = gate_mlp.unsqueeze(1) * ff_output452 453            hidden_states = ff_output + hidden_states454 455            return hidden_states456 457        def hacked_mid_forward(self, *args, **kwargs):458            eps = 1e-6459            x = self.original_forward(*args, **kwargs)460            if MODE == "write":461                if gn_auto_machine_weight >= self.gn_weight:462                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)463                    self.mean_bank.append(mean)464                    self.var_bank.append(var)465            if MODE == "read":466                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:467                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)468                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5469                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))470                    var_acc = sum(self.var_bank) / float(len(self.var_bank))471                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5472                    x_uc = (((x - mean) / std) * std_acc) + mean_acc473                    x_c = x_uc.clone()474                    if do_classifier_free_guidance and style_fidelity > 0:475                        x_c[uc_mask] = x[uc_mask]476                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc477                self.mean_bank = []478                self.var_bank = []479            return x480 481        def hack_CrossAttnDownBlock2D_forward(482            self,483            hidden_states: torch.FloatTensor,484            temb: Optional[torch.FloatTensor] = None,485            encoder_hidden_states: Optional[torch.FloatTensor] = None,486            attention_mask: Optional[torch.FloatTensor] = None,487            cross_attention_kwargs: Optional[Dict[str, Any]] = None,488            encoder_attention_mask: Optional[torch.FloatTensor] = None,489        ):490            eps = 1e-6491 492            # TODO(Patrick, William) - attention mask is not used493            output_states = ()494 495            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):496                hidden_states = resnet(hidden_states, temb)497                hidden_states = attn(498                    hidden_states,499                    encoder_hidden_states=encoder_hidden_states,500                    cross_attention_kwargs=cross_attention_kwargs,501                    attention_mask=attention_mask,502                    encoder_attention_mask=encoder_attention_mask,503                    return_dict=False,504                )[0]505                if MODE == "write":506                    if gn_auto_machine_weight >= self.gn_weight:507                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)508                        self.mean_bank.append([mean])509                        self.var_bank.append([var])510                if MODE == "read":511                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:512                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)513                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5514                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))515                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))516                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5517                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc518                        hidden_states_c = hidden_states_uc.clone()519                        if do_classifier_free_guidance and style_fidelity > 0:520                            hidden_states_c[uc_mask] = hidden_states[uc_mask]521                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc522 523                output_states = output_states + (hidden_states,)524 525            if MODE == "read":526                self.mean_bank = []527                self.var_bank = []528 529            if self.downsamplers is not None:530                for downsampler in self.downsamplers:531                    hidden_states = downsampler(hidden_states)532 533                output_states = output_states + (hidden_states,)534 535            return hidden_states, output_states536 537        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):538            eps = 1e-6539 540            output_states = ()541 542            for i, resnet in enumerate(self.resnets):543                hidden_states = resnet(hidden_states, temb)544 545                if MODE == "write":546                    if gn_auto_machine_weight >= self.gn_weight:547                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)548                        self.mean_bank.append([mean])549                        self.var_bank.append([var])550                if MODE == "read":551                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:552                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)553                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5554                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))555                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))556                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5557                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc558                        hidden_states_c = hidden_states_uc.clone()559                        if do_classifier_free_guidance and style_fidelity > 0:560                            hidden_states_c[uc_mask] = hidden_states[uc_mask]561                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc562 563                output_states = output_states + (hidden_states,)564 565            if MODE == "read":566                self.mean_bank = []567                self.var_bank = []568 569            if self.downsamplers is not None:570                for downsampler in self.downsamplers:571                    hidden_states = downsampler(hidden_states)572 573                output_states = output_states + (hidden_states,)574 575            return hidden_states, output_states576 577        def hacked_CrossAttnUpBlock2D_forward(578            self,579            hidden_states: torch.FloatTensor,580            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],581            temb: Optional[torch.FloatTensor] = None,582            encoder_hidden_states: Optional[torch.FloatTensor] = None,583            cross_attention_kwargs: Optional[Dict[str, Any]] = None,584            upsample_size: Optional[int] = None,585            attention_mask: Optional[torch.FloatTensor] = None,586            encoder_attention_mask: Optional[torch.FloatTensor] = None,587        ):588            eps = 1e-6589            # TODO(Patrick, William) - attention mask is not used590            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):591                # pop res hidden states592                res_hidden_states = res_hidden_states_tuple[-1]593                res_hidden_states_tuple = res_hidden_states_tuple[:-1]594                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)595                hidden_states = resnet(hidden_states, temb)596                hidden_states = attn(597                    hidden_states,598                    encoder_hidden_states=encoder_hidden_states,599                    cross_attention_kwargs=cross_attention_kwargs,600                    attention_mask=attention_mask,601                    encoder_attention_mask=encoder_attention_mask,602                    return_dict=False,603                )[0]604 605                if MODE == "write":606                    if gn_auto_machine_weight >= self.gn_weight:607                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)608                        self.mean_bank.append([mean])609                        self.var_bank.append([var])610                if MODE == "read":611                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:612                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)613                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5614                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))615                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))616                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5617                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc618                        hidden_states_c = hidden_states_uc.clone()619                        if do_classifier_free_guidance and style_fidelity > 0:620                            hidden_states_c[uc_mask] = hidden_states[uc_mask]621                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc622 623            if MODE == "read":624                self.mean_bank = []625                self.var_bank = []626 627            if self.upsamplers is not None:628                for upsampler in self.upsamplers:629                    hidden_states = upsampler(hidden_states, upsample_size)630 631            return hidden_states632 633        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):634            eps = 1e-6635            for i, resnet in enumerate(self.resnets):636                # pop res hidden states637                res_hidden_states = res_hidden_states_tuple[-1]638                res_hidden_states_tuple = res_hidden_states_tuple[:-1]639                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)640                hidden_states = resnet(hidden_states, temb)641 642                if MODE == "write":643                    if gn_auto_machine_weight >= self.gn_weight:644                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)645                        self.mean_bank.append([mean])646                        self.var_bank.append([var])647                if MODE == "read":648                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:649                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)650                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5651                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))652                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))653                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5654                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc655                        hidden_states_c = hidden_states_uc.clone()656                        if do_classifier_free_guidance and style_fidelity > 0:657                            hidden_states_c[uc_mask] = hidden_states[uc_mask]658                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc659 660            if MODE == "read":661                self.mean_bank = []662                self.var_bank = []663 664            if self.upsamplers is not None:665                for upsampler in self.upsamplers:666                    hidden_states = upsampler(hidden_states, upsample_size)667 668            return hidden_states669 670        if reference_attn:671            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]672            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])673 674            for i, module in enumerate(attn_modules):675                module._original_inner_forward = module.forward676                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)677                module.bank = []678                module.attn_weight = float(i) / float(len(attn_modules))679 680        if reference_adain:681            gn_modules = [self.unet.mid_block]682            self.unet.mid_block.gn_weight = 0683 684            down_blocks = self.unet.down_blocks685            for w, module in enumerate(down_blocks):686                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))687                gn_modules.append(module)688 689            up_blocks = self.unet.up_blocks690            for w, module in enumerate(up_blocks):691                module.gn_weight = float(w) / float(len(up_blocks))692                gn_modules.append(module)693 694            for i, module in enumerate(gn_modules):695                if getattr(module, "original_forward", None) is None:696                    module.original_forward = module.forward697                if i == 0:698                    # mid_block699                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)700                elif isinstance(module, CrossAttnDownBlock2D):701                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)702                elif isinstance(module, DownBlock2D):703                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)704                elif isinstance(module, CrossAttnUpBlock2D):705                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)706                elif isinstance(module, UpBlock2D):707                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)708                module.mean_bank = []709                module.var_bank = []710                module.gn_weight *= 2711 712        # 11. Denoising loop713        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order714        with self.progress_bar(total=num_inference_steps) as progress_bar:715            for i, t in enumerate(timesteps):716                # expand the latents if we are doing classifier free guidance717                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents718                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)719 720                # controlnet(s) inference721                if guess_mode and do_classifier_free_guidance:722                    # Infer ControlNet only for the conditional batch.723                    controlnet_latent_model_input = latents724                    controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]725                else:726                    controlnet_latent_model_input = latent_model_input727                    controlnet_prompt_embeds = prompt_embeds728 729                down_block_res_samples, mid_block_res_sample = self.controlnet(730                    controlnet_latent_model_input,731                    t,732                    encoder_hidden_states=controlnet_prompt_embeds,733                    controlnet_cond=image,734                    conditioning_scale=controlnet_conditioning_scale,735                    guess_mode=guess_mode,736                    return_dict=False,737                )738 739                if guess_mode and do_classifier_free_guidance:740                    # Infered ControlNet only for the conditional batch.741                    # To apply the output of ControlNet to both the unconditional and conditional batches,742                    # add 0 to the unconditional batch to keep it unchanged.743                    down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]744                    mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])745 746                # ref only part747                noise = randn_tensor(748                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype749                )750                ref_xt = self.scheduler.add_noise(751                    ref_image_latents,752                    noise,753                    t.reshape(754                        1,755                    ),756                )757                ref_xt = self.scheduler.scale_model_input(ref_xt, t)758 759                MODE = "write"760                self.unet(761                    ref_xt,762                    t,763                    encoder_hidden_states=prompt_embeds,764                    cross_attention_kwargs=cross_attention_kwargs,765                    return_dict=False,766                )767 768                # predict the noise residual769                MODE = "read"770                noise_pred = self.unet(771                    latent_model_input,772                    t,773                    encoder_hidden_states=prompt_embeds,774                    cross_attention_kwargs=cross_attention_kwargs,775                    down_block_additional_residuals=down_block_res_samples,776                    mid_block_additional_residual=mid_block_res_sample,777                    return_dict=False,778                )[0]779 780                # perform guidance781                if do_classifier_free_guidance:782                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)783                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)784 785                # compute the previous noisy sample x_t -> x_t-1786                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]787 788                # call the callback, if provided789                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):790                    progress_bar.update()791                    if callback is not None and i % callback_steps == 0:792                        callback(i, t, latents)793 794        # If we do sequential model offloading, let's offload unet and controlnet795        # manually for max memory savings796        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:797            self.unet.to("cpu")798            self.controlnet.to("cpu")799            torch.cuda.empty_cache()800 801        if not output_type == "latent":802            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]803            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)804        else:805            image = latents806            has_nsfw_concept = None807 808        if has_nsfw_concept is None:809            do_denormalize = [True] * image.shape[0]810        else:811            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]812 813        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)814 815        # Offload last model to CPU816        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:817            self.final_offload_hook.offload()818 819        if not return_dict:820            return (image, has_nsfw_concept)821 822        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)823