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 logging15from diffusers.utils.torch_utils import is_compiled_module, randn_tensor16 17 18logger = logging.get_logger(__name__)  # pylint: disable=invalid-name19 20EXAMPLE_DOC_STRING = """21    Examples:22        ```py23        >>> import cv224        >>> import torch25        >>> import numpy as np26        >>> from PIL import Image27        >>> from diffusers import UniPCMultistepScheduler28        >>> from diffusers.utils import load_image29 30        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")31 32        >>> # get canny image33        >>> image = cv2.Canny(np.array(input_image), 100, 200)34        >>> image = image[:, :, None]35        >>> image = np.concatenate([image, image, image], axis=2)36        >>> canny_image = Image.fromarray(image)37 38        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)39        >>> pipe = StableDiffusionControlNetReferencePipeline.from_pretrained(40                "runwayml/stable-diffusion-v1-5",41                controlnet=controlnet,42                safety_checker=None,43                torch_dtype=torch.float1644                ).to('cuda:0')45 46        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)47 48        >>> result_img = pipe(ref_image=input_image,49                        prompt="1girl",50                        image=canny_image,51                        num_inference_steps=20,52                        reference_attn=True,53                        reference_adain=True).images[0]54 55        >>> result_img.show()56        ```57"""58 59 60def torch_dfs(model: torch.nn.Module):61    result = [model]62    for child in model.children():63        result += torch_dfs(child)64    return result65 66 67class StableDiffusionControlNetReferencePipeline(StableDiffusionControlNetPipeline):68    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):69        refimage = refimage.to(device=device, dtype=dtype)70 71        # encode the mask image into latents space so we can concatenate it to the latents72        if isinstance(generator, list):73            ref_image_latents = [74                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])75                for i in range(batch_size)76            ]77            ref_image_latents = torch.cat(ref_image_latents, dim=0)78        else:79            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)80        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents81 82        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method83        if ref_image_latents.shape[0] < batch_size:84            if not batch_size % ref_image_latents.shape[0] == 0:85                raise ValueError(86                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"87                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."88                    " Make sure the number of images that you pass is divisible by the total requested batch size."89                )90            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)91 92        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents93 94        # aligning device to prevent device errors when concating it with the latent model input95        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)96        return ref_image_latents97 98    @torch.no_grad()99    def __call__(100        self,101        prompt: Union[str, List[str]] = None,102        image: Union[103            torch.FloatTensor,104            PIL.Image.Image,105            np.ndarray,106            List[torch.FloatTensor],107            List[PIL.Image.Image],108            List[np.ndarray],109        ] = None,110        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,111        height: Optional[int] = None,112        width: Optional[int] = None,113        num_inference_steps: int = 50,114        guidance_scale: float = 7.5,115        negative_prompt: Optional[Union[str, List[str]]] = None,116        num_images_per_prompt: Optional[int] = 1,117        eta: float = 0.0,118        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,119        latents: Optional[torch.FloatTensor] = None,120        prompt_embeds: Optional[torch.FloatTensor] = None,121        negative_prompt_embeds: Optional[torch.FloatTensor] = None,122        output_type: Optional[str] = "pil",123        return_dict: bool = True,124        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,125        callback_steps: int = 1,126        cross_attention_kwargs: Optional[Dict[str, Any]] = None,127        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,128        guess_mode: bool = False,129        attention_auto_machine_weight: float = 1.0,130        gn_auto_machine_weight: float = 1.0,131        style_fidelity: float = 0.5,132        reference_attn: bool = True,133        reference_adain: bool = True,134    ):135        r"""136        Function invoked when calling the pipeline for generation.137 138        Args:139            prompt (`str` or `List[str]`, *optional*):140                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.141                instead.142            image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:143                    `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):144                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If145                the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can146                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If147                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are148                specified in init, images must be passed as a list such that each element of the list can be correctly149                batched for input to a single controlnet.150            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):151                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If152                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can153                also be accepted as an image.154            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):155                The height in pixels of the generated image.156            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):157                The width in pixels of the generated image.158            num_inference_steps (`int`, *optional*, defaults to 50):159                The number of denoising steps. More denoising steps usually lead to a higher quality image at the160                expense of slower inference.161            guidance_scale (`float`, *optional*, defaults to 7.5):162                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).163                `guidance_scale` is defined as `w` of equation 2. of [Imagen164                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >165                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,166                usually at the expense of lower image quality.167            negative_prompt (`str` or `List[str]`, *optional*):168                The prompt or prompts not to guide the image generation. If not defined, one has to pass169                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is170                less than `1`).171            num_images_per_prompt (`int`, *optional*, defaults to 1):172                The number of images to generate per prompt.173            eta (`float`, *optional*, defaults to 0.0):174                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to175                [`schedulers.DDIMScheduler`], will be ignored for others.176            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):177                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)178                to make generation deterministic.179            latents (`torch.FloatTensor`, *optional*):180                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image181                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents182                tensor will ge generated by sampling using the supplied random `generator`.183            prompt_embeds (`torch.FloatTensor`, *optional*):184                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not185                provided, text embeddings will be generated from `prompt` input argument.186            negative_prompt_embeds (`torch.FloatTensor`, *optional*):187                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt188                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input189                argument.190            output_type (`str`, *optional*, defaults to `"pil"`):191                The output format of the generate image. Choose between192                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.193            return_dict (`bool`, *optional*, defaults to `True`):194                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a195                plain tuple.196            callback (`Callable`, *optional*):197                A function that will be called every `callback_steps` steps during inference. The function will be198                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.199            callback_steps (`int`, *optional*, defaults to 1):200                The frequency at which the `callback` function will be called. If not specified, the callback will be201                called at every step.202            cross_attention_kwargs (`dict`, *optional*):203                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under204                `self.processor` in205                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).206            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):207                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added208                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the209                corresponding scale as a list.210            guess_mode (`bool`, *optional*, defaults to `False`):211                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if212                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.213            attention_auto_machine_weight (`float`):214                Weight of using reference query for self attention's context.215                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.216            gn_auto_machine_weight (`float`):217                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.218            style_fidelity (`float`):219                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,220                elif style_fidelity=0.0, prompt more important, else balanced.221            reference_attn (`bool`):222                Whether to use reference query for self attention's context.223            reference_adain (`bool`):224                Whether to use reference adain.225 226        Examples:227 228        Returns:229            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:230            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.231            When returning a tuple, the first element is a list with the generated images, and the second element is a232            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"233            (nsfw) content, according to the `safety_checker`.234        """235        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."236 237        # 1. Check inputs. Raise error if not correct238        self.check_inputs(239            prompt,240            image,241            callback_steps,242            negative_prompt,243            prompt_embeds,244            negative_prompt_embeds,245            controlnet_conditioning_scale,246        )247 248        # 2. Define call parameters249        if prompt is not None and isinstance(prompt, str):250            batch_size = 1251        elif prompt is not None and isinstance(prompt, list):252            batch_size = len(prompt)253        else:254            batch_size = prompt_embeds.shape[0]255 256        device = self._execution_device257        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)258        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`259        # corresponds to doing no classifier free guidance.260        do_classifier_free_guidance = guidance_scale > 1.0261 262        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet263 264        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):265            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)266 267        global_pool_conditions = (268            controlnet.config.global_pool_conditions269            if isinstance(controlnet, ControlNetModel)270            else controlnet.nets[0].config.global_pool_conditions271        )272        guess_mode = guess_mode or global_pool_conditions273 274        # 3. Encode input prompt275        text_encoder_lora_scale = (276            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None277        )278        prompt_embeds = self._encode_prompt(279            prompt,280            device,281            num_images_per_prompt,282            do_classifier_free_guidance,283            negative_prompt,284            prompt_embeds=prompt_embeds,285            negative_prompt_embeds=negative_prompt_embeds,286            lora_scale=text_encoder_lora_scale,287        )288 289        # 4. Prepare image290        if isinstance(controlnet, ControlNetModel):291            image = self.prepare_image(292                image=image,293                width=width,294                height=height,295                batch_size=batch_size * num_images_per_prompt,296                num_images_per_prompt=num_images_per_prompt,297                device=device,298                dtype=controlnet.dtype,299                do_classifier_free_guidance=do_classifier_free_guidance,300                guess_mode=guess_mode,301            )302            height, width = image.shape[-2:]303        elif isinstance(controlnet, MultiControlNetModel):304            images = []305 306            for image_ in image:307                image_ = self.prepare_image(308                    image=image_,309                    width=width,310                    height=height,311                    batch_size=batch_size * num_images_per_prompt,312                    num_images_per_prompt=num_images_per_prompt,313                    device=device,314                    dtype=controlnet.dtype,315                    do_classifier_free_guidance=do_classifier_free_guidance,316                    guess_mode=guess_mode,317                )318 319                images.append(image_)320 321            image = images322            height, width = image[0].shape[-2:]323        else:324            assert False325 326        # 5. Preprocess reference image327        ref_image = self.prepare_image(328            image=ref_image,329            width=width,330            height=height,331            batch_size=batch_size * num_images_per_prompt,332            num_images_per_prompt=num_images_per_prompt,333            device=device,334            dtype=prompt_embeds.dtype,335        )336 337        # 6. Prepare timesteps338        self.scheduler.set_timesteps(num_inference_steps, device=device)339        timesteps = self.scheduler.timesteps340 341        # 7. Prepare latent variables342        num_channels_latents = self.unet.config.in_channels343        latents = self.prepare_latents(344            batch_size * num_images_per_prompt,345            num_channels_latents,346            height,347            width,348            prompt_embeds.dtype,349            device,350            generator,351            latents,352        )353 354        # 8. Prepare reference latent variables355        ref_image_latents = self.prepare_ref_latents(356            ref_image,357            batch_size * num_images_per_prompt,358            prompt_embeds.dtype,359            device,360            generator,361            do_classifier_free_guidance,362        )363 364        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline365        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)366 367        # 9. Modify self attention and group norm368        MODE = "write"369        uc_mask = (370            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)371            .type_as(ref_image_latents)372            .bool()373        )374 375        def hacked_basic_transformer_inner_forward(376            self,377            hidden_states: torch.FloatTensor,378            attention_mask: Optional[torch.FloatTensor] = None,379            encoder_hidden_states: Optional[torch.FloatTensor] = None,380            encoder_attention_mask: Optional[torch.FloatTensor] = None,381            timestep: Optional[torch.LongTensor] = None,382            cross_attention_kwargs: Dict[str, Any] = None,383            class_labels: Optional[torch.LongTensor] = None,384        ):385            if self.use_ada_layer_norm:386                norm_hidden_states = self.norm1(hidden_states, timestep)387            elif self.use_ada_layer_norm_zero:388                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(389                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype390                )391            else:392                norm_hidden_states = self.norm1(hidden_states)393 394            # 1. Self-Attention395            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}396            if self.only_cross_attention:397                attn_output = self.attn1(398                    norm_hidden_states,399                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,400                    attention_mask=attention_mask,401                    **cross_attention_kwargs,402                )403            else:404                if MODE == "write":405                    self.bank.append(norm_hidden_states.detach().clone())406                    attn_output = self.attn1(407                        norm_hidden_states,408                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,409                        attention_mask=attention_mask,410                        **cross_attention_kwargs,411                    )412                if MODE == "read":413                    if attention_auto_machine_weight > self.attn_weight:414                        attn_output_uc = self.attn1(415                            norm_hidden_states,416                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),417                            # attention_mask=attention_mask,418                            **cross_attention_kwargs,419                        )420                        attn_output_c = attn_output_uc.clone()421                        if do_classifier_free_guidance and style_fidelity > 0:422                            attn_output_c[uc_mask] = self.attn1(423                                norm_hidden_states[uc_mask],424                                encoder_hidden_states=norm_hidden_states[uc_mask],425                                **cross_attention_kwargs,426                            )427                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc428                        self.bank.clear()429                    else:430                        attn_output = self.attn1(431                            norm_hidden_states,432                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,433                            attention_mask=attention_mask,434                            **cross_attention_kwargs,435                        )436            if self.use_ada_layer_norm_zero:437                attn_output = gate_msa.unsqueeze(1) * attn_output438            hidden_states = attn_output + hidden_states439 440            if self.attn2 is not None:441                norm_hidden_states = (442                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)443                )444 445                # 2. Cross-Attention446                attn_output = self.attn2(447                    norm_hidden_states,448                    encoder_hidden_states=encoder_hidden_states,449                    attention_mask=encoder_attention_mask,450                    **cross_attention_kwargs,451                )452                hidden_states = attn_output + hidden_states453 454            # 3. Feed-forward455            norm_hidden_states = self.norm3(hidden_states)456 457            if self.use_ada_layer_norm_zero:458                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]459 460            ff_output = self.ff(norm_hidden_states)461 462            if self.use_ada_layer_norm_zero:463                ff_output = gate_mlp.unsqueeze(1) * ff_output464 465            hidden_states = ff_output + hidden_states466 467            return hidden_states468 469        def hacked_mid_forward(self, *args, **kwargs):470            eps = 1e-6471            x = self.original_forward(*args, **kwargs)472            if MODE == "write":473                if gn_auto_machine_weight >= self.gn_weight:474                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)475                    self.mean_bank.append(mean)476                    self.var_bank.append(var)477            if MODE == "read":478                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:479                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)480                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5481                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))482                    var_acc = sum(self.var_bank) / float(len(self.var_bank))483                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5484                    x_uc = (((x - mean) / std) * std_acc) + mean_acc485                    x_c = x_uc.clone()486                    if do_classifier_free_guidance and style_fidelity > 0:487                        x_c[uc_mask] = x[uc_mask]488                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc489                self.mean_bank = []490                self.var_bank = []491            return x492 493        def hack_CrossAttnDownBlock2D_forward(494            self,495            hidden_states: torch.FloatTensor,496            temb: Optional[torch.FloatTensor] = None,497            encoder_hidden_states: Optional[torch.FloatTensor] = None,498            attention_mask: Optional[torch.FloatTensor] = None,499            cross_attention_kwargs: Optional[Dict[str, Any]] = None,500            encoder_attention_mask: Optional[torch.FloatTensor] = None,501        ):502            eps = 1e-6503 504            # TODO(Patrick, William) - attention mask is not used505            output_states = ()506 507            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):508                hidden_states = resnet(hidden_states, temb)509                hidden_states = attn(510                    hidden_states,511                    encoder_hidden_states=encoder_hidden_states,512                    cross_attention_kwargs=cross_attention_kwargs,513                    attention_mask=attention_mask,514                    encoder_attention_mask=encoder_attention_mask,515                    return_dict=False,516                )[0]517                if MODE == "write":518                    if gn_auto_machine_weight >= self.gn_weight:519                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)520                        self.mean_bank.append([mean])521                        self.var_bank.append([var])522                if MODE == "read":523                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:524                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)525                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5526                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))527                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))528                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5529                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc530                        hidden_states_c = hidden_states_uc.clone()531                        if do_classifier_free_guidance and style_fidelity > 0:532                            hidden_states_c[uc_mask] = hidden_states[uc_mask]533                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc534 535                output_states = output_states + (hidden_states,)536 537            if MODE == "read":538                self.mean_bank = []539                self.var_bank = []540 541            if self.downsamplers is not None:542                for downsampler in self.downsamplers:543                    hidden_states = downsampler(hidden_states)544 545                output_states = output_states + (hidden_states,)546 547            return hidden_states, output_states548 549        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):550            eps = 1e-6551 552            output_states = ()553 554            for i, resnet in enumerate(self.resnets):555                hidden_states = resnet(hidden_states, temb)556 557                if MODE == "write":558                    if gn_auto_machine_weight >= self.gn_weight:559                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)560                        self.mean_bank.append([mean])561                        self.var_bank.append([var])562                if MODE == "read":563                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:564                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)565                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5566                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))567                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))568                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5569                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc570                        hidden_states_c = hidden_states_uc.clone()571                        if do_classifier_free_guidance and style_fidelity > 0:572                            hidden_states_c[uc_mask] = hidden_states[uc_mask]573                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc574 575                output_states = output_states + (hidden_states,)576 577            if MODE == "read":578                self.mean_bank = []579                self.var_bank = []580 581            if self.downsamplers is not None:582                for downsampler in self.downsamplers:583                    hidden_states = downsampler(hidden_states)584 585                output_states = output_states + (hidden_states,)586 587            return hidden_states, output_states588 589        def hacked_CrossAttnUpBlock2D_forward(590            self,591            hidden_states: torch.FloatTensor,592            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],593            temb: Optional[torch.FloatTensor] = None,594            encoder_hidden_states: Optional[torch.FloatTensor] = None,595            cross_attention_kwargs: Optional[Dict[str, Any]] = None,596            upsample_size: Optional[int] = None,597            attention_mask: Optional[torch.FloatTensor] = None,598            encoder_attention_mask: Optional[torch.FloatTensor] = None,599        ):600            eps = 1e-6601            # TODO(Patrick, William) - attention mask is not used602            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):603                # pop res hidden states604                res_hidden_states = res_hidden_states_tuple[-1]605                res_hidden_states_tuple = res_hidden_states_tuple[:-1]606                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)607                hidden_states = resnet(hidden_states, temb)608                hidden_states = attn(609                    hidden_states,610                    encoder_hidden_states=encoder_hidden_states,611                    cross_attention_kwargs=cross_attention_kwargs,612                    attention_mask=attention_mask,613                    encoder_attention_mask=encoder_attention_mask,614                    return_dict=False,615                )[0]616 617                if MODE == "write":618                    if gn_auto_machine_weight >= self.gn_weight:619                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)620                        self.mean_bank.append([mean])621                        self.var_bank.append([var])622                if MODE == "read":623                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:624                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)625                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5626                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))627                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))628                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5629                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc630                        hidden_states_c = hidden_states_uc.clone()631                        if do_classifier_free_guidance and style_fidelity > 0:632                            hidden_states_c[uc_mask] = hidden_states[uc_mask]633                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc634 635            if MODE == "read":636                self.mean_bank = []637                self.var_bank = []638 639            if self.upsamplers is not None:640                for upsampler in self.upsamplers:641                    hidden_states = upsampler(hidden_states, upsample_size)642 643            return hidden_states644 645        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):646            eps = 1e-6647            for i, resnet in enumerate(self.resnets):648                # pop res hidden states649                res_hidden_states = res_hidden_states_tuple[-1]650                res_hidden_states_tuple = res_hidden_states_tuple[:-1]651                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)652                hidden_states = resnet(hidden_states, temb)653 654                if MODE == "write":655                    if gn_auto_machine_weight >= self.gn_weight:656                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)657                        self.mean_bank.append([mean])658                        self.var_bank.append([var])659                if MODE == "read":660                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:661                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)662                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5663                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))664                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))665                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5666                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc667                        hidden_states_c = hidden_states_uc.clone()668                        if do_classifier_free_guidance and style_fidelity > 0:669                            hidden_states_c[uc_mask] = hidden_states[uc_mask]670                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc671 672            if MODE == "read":673                self.mean_bank = []674                self.var_bank = []675 676            if self.upsamplers is not None:677                for upsampler in self.upsamplers:678                    hidden_states = upsampler(hidden_states, upsample_size)679 680            return hidden_states681 682        if reference_attn:683            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]684            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])685 686            for i, module in enumerate(attn_modules):687                module._original_inner_forward = module.forward688                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)689                module.bank = []690                module.attn_weight = float(i) / float(len(attn_modules))691 692        if reference_adain:693            gn_modules = [self.unet.mid_block]694            self.unet.mid_block.gn_weight = 0695 696            down_blocks = self.unet.down_blocks697            for w, module in enumerate(down_blocks):698                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))699                gn_modules.append(module)700 701            up_blocks = self.unet.up_blocks702            for w, module in enumerate(up_blocks):703                module.gn_weight = float(w) / float(len(up_blocks))704                gn_modules.append(module)705 706            for i, module in enumerate(gn_modules):707                if getattr(module, "original_forward", None) is None:708                    module.original_forward = module.forward709                if i == 0:710                    # mid_block711                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)712                elif isinstance(module, CrossAttnDownBlock2D):713                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)714                elif isinstance(module, DownBlock2D):715                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)716                elif isinstance(module, CrossAttnUpBlock2D):717                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)718                elif isinstance(module, UpBlock2D):719                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)720                module.mean_bank = []721                module.var_bank = []722                module.gn_weight *= 2723 724        # 11. Denoising loop725        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order726        with self.progress_bar(total=num_inference_steps) as progress_bar:727            for i, t in enumerate(timesteps):728                # expand the latents if we are doing classifier free guidance729                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents730                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)731 732                # controlnet(s) inference733                if guess_mode and do_classifier_free_guidance:734                    # Infer ControlNet only for the conditional batch.735                    control_model_input = latents736                    control_model_input = self.scheduler.scale_model_input(control_model_input, t)737                    controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]738                else:739                    control_model_input = latent_model_input740                    controlnet_prompt_embeds = prompt_embeds741 742                down_block_res_samples, mid_block_res_sample = self.controlnet(743                    control_model_input,744                    t,745                    encoder_hidden_states=controlnet_prompt_embeds,746                    controlnet_cond=image,747                    conditioning_scale=controlnet_conditioning_scale,748                    guess_mode=guess_mode,749                    return_dict=False,750                )751 752                if guess_mode and do_classifier_free_guidance:753                    # Infered ControlNet only for the conditional batch.754                    # To apply the output of ControlNet to both the unconditional and conditional batches,755                    # add 0 to the unconditional batch to keep it unchanged.756                    down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]757                    mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])758 759                # ref only part760                noise = randn_tensor(761                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype762                )763                ref_xt = self.scheduler.add_noise(764                    ref_image_latents,765                    noise,766                    t.reshape(767                        1,768                    ),769                )770                ref_xt = self.scheduler.scale_model_input(ref_xt, t)771 772                MODE = "write"773                self.unet(774                    ref_xt,775                    t,776                    encoder_hidden_states=prompt_embeds,777                    cross_attention_kwargs=cross_attention_kwargs,778                    return_dict=False,779                )780 781                # predict the noise residual782                MODE = "read"783                noise_pred = self.unet(784                    latent_model_input,785                    t,786                    encoder_hidden_states=prompt_embeds,787                    cross_attention_kwargs=cross_attention_kwargs,788                    down_block_additional_residuals=down_block_res_samples,789                    mid_block_additional_residual=mid_block_res_sample,790                    return_dict=False,791                )[0]792 793                # perform guidance794                if do_classifier_free_guidance:795                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)796                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)797 798                # compute the previous noisy sample x_t -> x_t-1799                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]800 801                # call the callback, if provided802                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):803                    progress_bar.update()804                    if callback is not None and i % callback_steps == 0:805                        step_idx = i // getattr(self.scheduler, "order", 1)806                        callback(step_idx, t, latents)807 808        # If we do sequential model offloading, let's offload unet and controlnet809        # manually for max memory savings810        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:811            self.unet.to("cpu")812            self.controlnet.to("cpu")813            torch.cuda.empty_cache()814 815        if not output_type == "latent":816            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]817            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)818        else:819            image = latents820            has_nsfw_concept = None821 822        if has_nsfw_concept is None:823            do_denormalize = [True] * image.shape[0]824        else:825            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]826 827        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)828 829        # Offload last model to CPU830        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:831            self.final_offload_hook.offload()832 833        if not return_dict:834            return (image, has_nsfw_concept)835 836        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)837