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 StableDiffusionPipeline9from diffusers.models.attention import BasicTransformerBlock10from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D11from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput12from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg13from diffusers.utils import PIL_INTERPOLATION, logging, randn_tensor14 15 16logger = logging.get_logger(__name__)  # pylint: disable=invalid-name17 18EXAMPLE_DOC_STRING = """19    Examples:20        ```py21        >>> import torch22        >>> from diffusers import UniPCMultistepScheduler23        >>> from diffusers.utils import load_image24 25        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")26 27        >>> pipe = StableDiffusionReferencePipeline.from_pretrained(28                "runwayml/stable-diffusion-v1-5",29                safety_checker=None,30                torch_dtype=torch.float1631                ).to('cuda:0')32 33        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)34 35        >>> result_img = pipe(ref_image=input_image,36                        prompt="1girl",37                        num_inference_steps=20,38                        reference_attn=True,39                        reference_adain=True).images[0]40 41        >>> result_img.show()42        ```43"""44 45 46def torch_dfs(model: torch.nn.Module):47    result = [model]48    for child in model.children():49        result += torch_dfs(child)50    return result51 52 53class StableDiffusionReferencePipeline(StableDiffusionPipeline):54    def _default_height_width(self, height, width, image):55        # NOTE: It is possible that a list of images have different56        # dimensions for each image, so just checking the first image57        # is not _exactly_ correct, but it is simple.58        while isinstance(image, list):59            image = image[0]60 61        if height is None:62            if isinstance(image, PIL.Image.Image):63                height = image.height64            elif isinstance(image, torch.Tensor):65                height = image.shape[2]66 67            height = (height // 8) * 8  # round down to nearest multiple of 868 69        if width is None:70            if isinstance(image, PIL.Image.Image):71                width = image.width72            elif isinstance(image, torch.Tensor):73                width = image.shape[3]74 75            width = (width // 8) * 8  # round down to nearest multiple of 876 77        return height, width78 79    def prepare_image(80        self,81        image,82        width,83        height,84        batch_size,85        num_images_per_prompt,86        device,87        dtype,88        do_classifier_free_guidance=False,89        guess_mode=False,90    ):91        if not isinstance(image, torch.Tensor):92            if isinstance(image, PIL.Image.Image):93                image = [image]94 95            if isinstance(image[0], PIL.Image.Image):96                images = []97 98                for image_ in image:99                    image_ = image_.convert("RGB")100                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])101                    image_ = np.array(image_)102                    image_ = image_[None, :]103                    images.append(image_)104 105                image = images106 107                image = np.concatenate(image, axis=0)108                image = np.array(image).astype(np.float32) / 255.0109                image = (image - 0.5) / 0.5110                image = image.transpose(0, 3, 1, 2)111                image = torch.from_numpy(image)112            elif isinstance(image[0], torch.Tensor):113                image = torch.cat(image, dim=0)114 115        image_batch_size = image.shape[0]116 117        if image_batch_size == 1:118            repeat_by = batch_size119        else:120            # image batch size is the same as prompt batch size121            repeat_by = num_images_per_prompt122 123        image = image.repeat_interleave(repeat_by, dim=0)124 125        image = image.to(device=device, dtype=dtype)126 127        if do_classifier_free_guidance and not guess_mode:128            image = torch.cat([image] * 2)129 130        return image131 132    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):133        refimage = refimage.to(device=device, dtype=dtype)134 135        # encode the mask image into latents space so we can concatenate it to the latents136        if isinstance(generator, list):137            ref_image_latents = [138                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])139                for i in range(batch_size)140            ]141            ref_image_latents = torch.cat(ref_image_latents, dim=0)142        else:143            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)144        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents145 146        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method147        if ref_image_latents.shape[0] < batch_size:148            if not batch_size % ref_image_latents.shape[0] == 0:149                raise ValueError(150                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"151                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."152                    " Make sure the number of images that you pass is divisible by the total requested batch size."153                )154            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)155 156        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents157 158        # aligning device to prevent device errors when concating it with the latent model input159        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)160        return ref_image_latents161 162    @torch.no_grad()163    def __call__(164        self,165        prompt: Union[str, List[str]] = None,166        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,167        height: Optional[int] = None,168        width: Optional[int] = None,169        num_inference_steps: int = 50,170        guidance_scale: float = 7.5,171        negative_prompt: Optional[Union[str, List[str]]] = None,172        num_images_per_prompt: Optional[int] = 1,173        eta: float = 0.0,174        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,175        latents: Optional[torch.FloatTensor] = None,176        prompt_embeds: Optional[torch.FloatTensor] = None,177        negative_prompt_embeds: Optional[torch.FloatTensor] = None,178        output_type: Optional[str] = "pil",179        return_dict: bool = True,180        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,181        callback_steps: int = 1,182        cross_attention_kwargs: Optional[Dict[str, Any]] = None,183        guidance_rescale: float = 0.0,184        attention_auto_machine_weight: float = 1.0,185        gn_auto_machine_weight: float = 1.0,186        style_fidelity: float = 0.5,187        reference_attn: bool = True,188        reference_adain: bool = True,189    ):190        r"""191        Function invoked when calling the pipeline for generation.192 193        Args:194            prompt (`str` or `List[str]`, *optional*):195                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.196                instead.197            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):198                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If199                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can200                also be accepted as an image.201            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):202                The height in pixels of the generated image.203            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):204                The width in pixels of the generated image.205            num_inference_steps (`int`, *optional*, defaults to 50):206                The number of denoising steps. More denoising steps usually lead to a higher quality image at the207                expense of slower inference.208            guidance_scale (`float`, *optional*, defaults to 7.5):209                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).210                `guidance_scale` is defined as `w` of equation 2. of [Imagen211                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >212                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,213                usually at the expense of lower image quality.214            negative_prompt (`str` or `List[str]`, *optional*):215                The prompt or prompts not to guide the image generation. If not defined, one has to pass216                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is217                less than `1`).218            num_images_per_prompt (`int`, *optional*, defaults to 1):219                The number of images to generate per prompt.220            eta (`float`, *optional*, defaults to 0.0):221                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to222                [`schedulers.DDIMScheduler`], will be ignored for others.223            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):224                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)225                to make generation deterministic.226            latents (`torch.FloatTensor`, *optional*):227                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image228                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents229                tensor will ge generated by sampling using the supplied random `generator`.230            prompt_embeds (`torch.FloatTensor`, *optional*):231                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not232                provided, text embeddings will be generated from `prompt` input argument.233            negative_prompt_embeds (`torch.FloatTensor`, *optional*):234                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt235                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input236                argument.237            output_type (`str`, *optional*, defaults to `"pil"`):238                The output format of the generate image. Choose between239                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.240            return_dict (`bool`, *optional*, defaults to `True`):241                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a242                plain tuple.243            callback (`Callable`, *optional*):244                A function that will be called every `callback_steps` steps during inference. The function will be245                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.246            callback_steps (`int`, *optional*, defaults to 1):247                The frequency at which the `callback` function will be called. If not specified, the callback will be248                called at every step.249            cross_attention_kwargs (`dict`, *optional*):250                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under251                `self.processor` in252                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).253            guidance_rescale (`float`, *optional*, defaults to 0.7):254                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are255                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of256                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).257                Guidance rescale factor should fix overexposure when using zero terminal SNR.258            attention_auto_machine_weight (`float`):259                Weight of using reference query for self attention's context.260                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.261            gn_auto_machine_weight (`float`):262                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.263            style_fidelity (`float`):264                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,265                elif style_fidelity=0.0, prompt more important, else balanced.266            reference_attn (`bool`):267                Whether to use reference query for self attention's context.268            reference_adain (`bool`):269                Whether to use reference adain.270 271        Examples:272 273        Returns:274            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:275            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.276            When returning a tuple, the first element is a list with the generated images, and the second element is a277            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"278            (nsfw) content, according to the `safety_checker`.279        """280        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."281 282        # 0. Default height and width to unet283        height, width = self._default_height_width(height, width, ref_image)284 285        # 1. Check inputs. Raise error if not correct286        self.check_inputs(287            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds288        )289 290        # 2. Define call parameters291        if prompt is not None and isinstance(prompt, str):292            batch_size = 1293        elif prompt is not None and isinstance(prompt, list):294            batch_size = len(prompt)295        else:296            batch_size = prompt_embeds.shape[0]297 298        device = self._execution_device299        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)300        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`301        # corresponds to doing no classifier free guidance.302        do_classifier_free_guidance = guidance_scale > 1.0303 304        # 3. Encode input prompt305        text_encoder_lora_scale = (306            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None307        )308        prompt_embeds = self._encode_prompt(309            prompt,310            device,311            num_images_per_prompt,312            do_classifier_free_guidance,313            negative_prompt,314            prompt_embeds=prompt_embeds,315            negative_prompt_embeds=negative_prompt_embeds,316            lora_scale=text_encoder_lora_scale,317        )318 319        # 4. Preprocess reference image320        ref_image = self.prepare_image(321            image=ref_image,322            width=width,323            height=height,324            batch_size=batch_size * num_images_per_prompt,325            num_images_per_prompt=num_images_per_prompt,326            device=device,327            dtype=prompt_embeds.dtype,328        )329 330        # 5. Prepare timesteps331        self.scheduler.set_timesteps(num_inference_steps, device=device)332        timesteps = self.scheduler.timesteps333 334        # 6. Prepare latent variables335        num_channels_latents = self.unet.config.in_channels336        latents = self.prepare_latents(337            batch_size * num_images_per_prompt,338            num_channels_latents,339            height,340            width,341            prompt_embeds.dtype,342            device,343            generator,344            latents,345        )346 347        # 7. Prepare reference latent variables348        ref_image_latents = self.prepare_ref_latents(349            ref_image,350            batch_size * num_images_per_prompt,351            prompt_embeds.dtype,352            device,353            generator,354            do_classifier_free_guidance,355        )356 357        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline358        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)359 360        # 9. Modify self attention and group norm361        MODE = "write"362        uc_mask = (363            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)364            .type_as(ref_image_latents)365            .bool()366        )367 368        def hacked_basic_transformer_inner_forward(369            self,370            hidden_states: torch.FloatTensor,371            attention_mask: Optional[torch.FloatTensor] = None,372            encoder_hidden_states: Optional[torch.FloatTensor] = None,373            encoder_attention_mask: Optional[torch.FloatTensor] = None,374            timestep: Optional[torch.LongTensor] = None,375            cross_attention_kwargs: Dict[str, Any] = None,376            class_labels: Optional[torch.LongTensor] = None,377        ):378            if self.use_ada_layer_norm:379                norm_hidden_states = self.norm1(hidden_states, timestep)380            elif self.use_ada_layer_norm_zero:381                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(382                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype383                )384            else:385                norm_hidden_states = self.norm1(hidden_states)386 387            # 1. Self-Attention388            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}389            if self.only_cross_attention:390                attn_output = self.attn1(391                    norm_hidden_states,392                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,393                    attention_mask=attention_mask,394                    **cross_attention_kwargs,395                )396            else:397                if MODE == "write":398                    self.bank.append(norm_hidden_states.detach().clone())399                    attn_output = self.attn1(400                        norm_hidden_states,401                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,402                        attention_mask=attention_mask,403                        **cross_attention_kwargs,404                    )405                if MODE == "read":406                    if attention_auto_machine_weight > self.attn_weight:407                        attn_output_uc = self.attn1(408                            norm_hidden_states,409                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),410                            # attention_mask=attention_mask,411                            **cross_attention_kwargs,412                        )413                        attn_output_c = attn_output_uc.clone()414                        if do_classifier_free_guidance and style_fidelity > 0:415                            attn_output_c[uc_mask] = self.attn1(416                                norm_hidden_states[uc_mask],417                                encoder_hidden_states=norm_hidden_states[uc_mask],418                                **cross_attention_kwargs,419                            )420                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc421                        self.bank.clear()422                    else:423                        attn_output = self.attn1(424                            norm_hidden_states,425                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,426                            attention_mask=attention_mask,427                            **cross_attention_kwargs,428                        )429            if self.use_ada_layer_norm_zero:430                attn_output = gate_msa.unsqueeze(1) * attn_output431            hidden_states = attn_output + hidden_states432 433            if self.attn2 is not None:434                norm_hidden_states = (435                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)436                )437 438                # 2. Cross-Attention439                attn_output = self.attn2(440                    norm_hidden_states,441                    encoder_hidden_states=encoder_hidden_states,442                    attention_mask=encoder_attention_mask,443                    **cross_attention_kwargs,444                )445                hidden_states = attn_output + hidden_states446 447            # 3. Feed-forward448            norm_hidden_states = self.norm3(hidden_states)449 450            if self.use_ada_layer_norm_zero:451                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]452 453            ff_output = self.ff(norm_hidden_states)454 455            if self.use_ada_layer_norm_zero:456                ff_output = gate_mlp.unsqueeze(1) * ff_output457 458            hidden_states = ff_output + hidden_states459 460            return hidden_states461 462        def hacked_mid_forward(self, *args, **kwargs):463            eps = 1e-6464            x = self.original_forward(*args, **kwargs)465            if MODE == "write":466                if gn_auto_machine_weight >= self.gn_weight:467                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)468                    self.mean_bank.append(mean)469                    self.var_bank.append(var)470            if MODE == "read":471                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:472                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)473                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5474                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))475                    var_acc = sum(self.var_bank) / float(len(self.var_bank))476                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5477                    x_uc = (((x - mean) / std) * std_acc) + mean_acc478                    x_c = x_uc.clone()479                    if do_classifier_free_guidance and style_fidelity > 0:480                        x_c[uc_mask] = x[uc_mask]481                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc482                self.mean_bank = []483                self.var_bank = []484            return x485 486        def hack_CrossAttnDownBlock2D_forward(487            self,488            hidden_states: torch.FloatTensor,489            temb: Optional[torch.FloatTensor] = None,490            encoder_hidden_states: Optional[torch.FloatTensor] = None,491            attention_mask: Optional[torch.FloatTensor] = None,492            cross_attention_kwargs: Optional[Dict[str, Any]] = None,493            encoder_attention_mask: Optional[torch.FloatTensor] = None,494        ):495            eps = 1e-6496 497            # TODO(Patrick, William) - attention mask is not used498            output_states = ()499 500            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):501                hidden_states = resnet(hidden_states, temb)502                hidden_states = attn(503                    hidden_states,504                    encoder_hidden_states=encoder_hidden_states,505                    cross_attention_kwargs=cross_attention_kwargs,506                    attention_mask=attention_mask,507                    encoder_attention_mask=encoder_attention_mask,508                    return_dict=False,509                )[0]510                if MODE == "write":511                    if gn_auto_machine_weight >= self.gn_weight:512                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)513                        self.mean_bank.append([mean])514                        self.var_bank.append([var])515                if MODE == "read":516                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:517                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)518                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5519                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))520                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))521                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5522                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc523                        hidden_states_c = hidden_states_uc.clone()524                        if do_classifier_free_guidance and style_fidelity > 0:525                            hidden_states_c[uc_mask] = hidden_states[uc_mask]526                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc527 528                output_states = output_states + (hidden_states,)529 530            if MODE == "read":531                self.mean_bank = []532                self.var_bank = []533 534            if self.downsamplers is not None:535                for downsampler in self.downsamplers:536                    hidden_states = downsampler(hidden_states)537 538                output_states = output_states + (hidden_states,)539 540            return hidden_states, output_states541 542        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):543            eps = 1e-6544 545            output_states = ()546 547            for i, resnet in enumerate(self.resnets):548                hidden_states = resnet(hidden_states, temb)549 550                if MODE == "write":551                    if gn_auto_machine_weight >= self.gn_weight:552                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)553                        self.mean_bank.append([mean])554                        self.var_bank.append([var])555                if MODE == "read":556                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:557                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)558                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5559                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))560                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))561                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5562                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc563                        hidden_states_c = hidden_states_uc.clone()564                        if do_classifier_free_guidance and style_fidelity > 0:565                            hidden_states_c[uc_mask] = hidden_states[uc_mask]566                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc567 568                output_states = output_states + (hidden_states,)569 570            if MODE == "read":571                self.mean_bank = []572                self.var_bank = []573 574            if self.downsamplers is not None:575                for downsampler in self.downsamplers:576                    hidden_states = downsampler(hidden_states)577 578                output_states = output_states + (hidden_states,)579 580            return hidden_states, output_states581 582        def hacked_CrossAttnUpBlock2D_forward(583            self,584            hidden_states: torch.FloatTensor,585            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],586            temb: Optional[torch.FloatTensor] = None,587            encoder_hidden_states: Optional[torch.FloatTensor] = None,588            cross_attention_kwargs: Optional[Dict[str, Any]] = None,589            upsample_size: Optional[int] = None,590            attention_mask: Optional[torch.FloatTensor] = None,591            encoder_attention_mask: Optional[torch.FloatTensor] = None,592        ):593            eps = 1e-6594            # TODO(Patrick, William) - attention mask is not used595            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):596                # pop res hidden states597                res_hidden_states = res_hidden_states_tuple[-1]598                res_hidden_states_tuple = res_hidden_states_tuple[:-1]599                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)600                hidden_states = resnet(hidden_states, temb)601                hidden_states = attn(602                    hidden_states,603                    encoder_hidden_states=encoder_hidden_states,604                    cross_attention_kwargs=cross_attention_kwargs,605                    attention_mask=attention_mask,606                    encoder_attention_mask=encoder_attention_mask,607                    return_dict=False,608                )[0]609 610                if MODE == "write":611                    if gn_auto_machine_weight >= self.gn_weight:612                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)613                        self.mean_bank.append([mean])614                        self.var_bank.append([var])615                if MODE == "read":616                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:617                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)618                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5619                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))620                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))621                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5622                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc623                        hidden_states_c = hidden_states_uc.clone()624                        if do_classifier_free_guidance and style_fidelity > 0:625                            hidden_states_c[uc_mask] = hidden_states[uc_mask]626                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc627 628            if MODE == "read":629                self.mean_bank = []630                self.var_bank = []631 632            if self.upsamplers is not None:633                for upsampler in self.upsamplers:634                    hidden_states = upsampler(hidden_states, upsample_size)635 636            return hidden_states637 638        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):639            eps = 1e-6640            for i, resnet in enumerate(self.resnets):641                # pop res hidden states642                res_hidden_states = res_hidden_states_tuple[-1]643                res_hidden_states_tuple = res_hidden_states_tuple[:-1]644                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)645                hidden_states = resnet(hidden_states, temb)646 647                if MODE == "write":648                    if gn_auto_machine_weight >= self.gn_weight:649                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)650                        self.mean_bank.append([mean])651                        self.var_bank.append([var])652                if MODE == "read":653                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:654                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)655                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5656                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))657                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))658                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5659                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc660                        hidden_states_c = hidden_states_uc.clone()661                        if do_classifier_free_guidance and style_fidelity > 0:662                            hidden_states_c[uc_mask] = hidden_states[uc_mask]663                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc664 665            if MODE == "read":666                self.mean_bank = []667                self.var_bank = []668 669            if self.upsamplers is not None:670                for upsampler in self.upsamplers:671                    hidden_states = upsampler(hidden_states, upsample_size)672 673            return hidden_states674 675        if reference_attn:676            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]677            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])678 679            for i, module in enumerate(attn_modules):680                module._original_inner_forward = module.forward681                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)682                module.bank = []683                module.attn_weight = float(i) / float(len(attn_modules))684 685        if reference_adain:686            gn_modules = [self.unet.mid_block]687            self.unet.mid_block.gn_weight = 0688 689            down_blocks = self.unet.down_blocks690            for w, module in enumerate(down_blocks):691                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))692                gn_modules.append(module)693 694            up_blocks = self.unet.up_blocks695            for w, module in enumerate(up_blocks):696                module.gn_weight = float(w) / float(len(up_blocks))697                gn_modules.append(module)698 699            for i, module in enumerate(gn_modules):700                if getattr(module, "original_forward", None) is None:701                    module.original_forward = module.forward702                if i == 0:703                    # mid_block704                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)705                elif isinstance(module, CrossAttnDownBlock2D):706                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)707                elif isinstance(module, DownBlock2D):708                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)709                elif isinstance(module, CrossAttnUpBlock2D):710                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)711                elif isinstance(module, UpBlock2D):712                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)713                module.mean_bank = []714                module.var_bank = []715                module.gn_weight *= 2716 717        # 10. Denoising loop718        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order719        with self.progress_bar(total=num_inference_steps) as progress_bar:720            for i, t in enumerate(timesteps):721                # expand the latents if we are doing classifier free guidance722                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents723                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)724 725                # ref only part726                noise = randn_tensor(727                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype728                )729                ref_xt = self.scheduler.add_noise(730                    ref_image_latents,731                    noise,732                    t.reshape(733                        1,734                    ),735                )736                ref_xt = self.scheduler.scale_model_input(ref_xt, t)737 738                MODE = "write"739                self.unet(740                    ref_xt,741                    t,742                    encoder_hidden_states=prompt_embeds,743                    cross_attention_kwargs=cross_attention_kwargs,744                    return_dict=False,745                )746 747                # predict the noise residual748                MODE = "read"749                noise_pred = self.unet(750                    latent_model_input,751                    t,752                    encoder_hidden_states=prompt_embeds,753                    cross_attention_kwargs=cross_attention_kwargs,754                    return_dict=False,755                )[0]756 757                # perform guidance758                if do_classifier_free_guidance:759                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)760                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)761 762                if do_classifier_free_guidance and guidance_rescale > 0.0:763                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf764                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)765 766                # compute the previous noisy sample x_t -> x_t-1767                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]768 769                # call the callback, if provided770                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):771                    progress_bar.update()772                    if callback is not None and i % callback_steps == 0:773                        callback(i, t, latents)774 775        if not output_type == "latent":776            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]777            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)778        else:779            image = latents780            has_nsfw_concept = None781 782        if has_nsfw_concept is None:783            do_denormalize = [True] * image.shape[0]784        else:785            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]786 787        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)788 789        # Offload last model to CPU790        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:791            self.final_offload_hook.offload()792 793        if not return_dict:794            return (image, has_nsfw_concept)795 796        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)797