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, logging14from diffusers.utils.torch_utils import randn_tensor15 16 17logger = logging.get_logger(__name__)  # pylint: disable=invalid-name18 19EXAMPLE_DOC_STRING = """20    Examples:21        ```py22        >>> import torch23        >>> from diffusers import UniPCMultistepScheduler24        >>> from diffusers.utils import load_image25 26        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")27 28        >>> pipe = StableDiffusionReferencePipeline.from_pretrained(29                "runwayml/stable-diffusion-v1-5",30                safety_checker=None,31                torch_dtype=torch.float1632                ).to('cuda:0')33 34        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)35 36        >>> result_img = pipe(ref_image=input_image,37                        prompt="1girl",38                        num_inference_steps=20,39                        reference_attn=True,40                        reference_adain=True).images[0]41 42        >>> result_img.show()43        ```44"""45 46 47def torch_dfs(model: torch.nn.Module):48    result = [model]49    for child in model.children():50        result += torch_dfs(child)51    return result52 53 54class StableDiffusionReferencePipeline(StableDiffusionPipeline):55    def _default_height_width(self, height, width, image):56        # NOTE: It is possible that a list of images have different57        # dimensions for each image, so just checking the first image58        # is not _exactly_ correct, but it is simple.59        while isinstance(image, list):60            image = image[0]61 62        if height is None:63            if isinstance(image, PIL.Image.Image):64                height = image.height65            elif isinstance(image, torch.Tensor):66                height = image.shape[2]67 68            height = (height // 8) * 8  # round down to nearest multiple of 869 70        if width is None:71            if isinstance(image, PIL.Image.Image):72                width = image.width73            elif isinstance(image, torch.Tensor):74                width = image.shape[3]75 76            width = (width // 8) * 8  # round down to nearest multiple of 877 78        return height, width79 80    def prepare_image(81        self,82        image,83        width,84        height,85        batch_size,86        num_images_per_prompt,87        device,88        dtype,89        do_classifier_free_guidance=False,90        guess_mode=False,91    ):92        if not isinstance(image, torch.Tensor):93            if isinstance(image, PIL.Image.Image):94                image = [image]95 96            if isinstance(image[0], PIL.Image.Image):97                images = []98 99                for image_ in image:100                    image_ = image_.convert("RGB")101                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])102                    image_ = np.array(image_)103                    image_ = image_[None, :]104                    images.append(image_)105 106                image = images107 108                image = np.concatenate(image, axis=0)109                image = np.array(image).astype(np.float32) / 255.0110                image = (image - 0.5) / 0.5111                image = image.transpose(0, 3, 1, 2)112                image = torch.from_numpy(image)113            elif isinstance(image[0], torch.Tensor):114                image = torch.cat(image, dim=0)115 116        image_batch_size = image.shape[0]117 118        if image_batch_size == 1:119            repeat_by = batch_size120        else:121            # image batch size is the same as prompt batch size122            repeat_by = num_images_per_prompt123 124        image = image.repeat_interleave(repeat_by, dim=0)125 126        image = image.to(device=device, dtype=dtype)127 128        if do_classifier_free_guidance and not guess_mode:129            image = torch.cat([image] * 2)130 131        return image132 133    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):134        refimage = refimage.to(device=device, dtype=dtype)135 136        # encode the mask image into latents space so we can concatenate it to the latents137        if isinstance(generator, list):138            ref_image_latents = [139                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])140                for i in range(batch_size)141            ]142            ref_image_latents = torch.cat(ref_image_latents, dim=0)143        else:144            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)145        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents146 147        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method148        if ref_image_latents.shape[0] < batch_size:149            if not batch_size % ref_image_latents.shape[0] == 0:150                raise ValueError(151                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"152                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."153                    " Make sure the number of images that you pass is divisible by the total requested batch size."154                )155            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)156 157        # aligning device to prevent device errors when concating it with the latent model input158        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)159        return ref_image_latents160 161    @torch.no_grad()162    def __call__(163        self,164        prompt: Union[str, List[str]] = None,165        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,166        height: Optional[int] = None,167        width: Optional[int] = None,168        num_inference_steps: int = 50,169        guidance_scale: float = 7.5,170        negative_prompt: Optional[Union[str, List[str]]] = None,171        num_images_per_prompt: Optional[int] = 1,172        eta: float = 0.0,173        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,174        latents: Optional[torch.FloatTensor] = None,175        prompt_embeds: Optional[torch.FloatTensor] = None,176        negative_prompt_embeds: Optional[torch.FloatTensor] = None,177        output_type: Optional[str] = "pil",178        return_dict: bool = True,179        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,180        callback_steps: int = 1,181        cross_attention_kwargs: Optional[Dict[str, Any]] = None,182        guidance_rescale: float = 0.0,183        attention_auto_machine_weight: float = 1.0,184        gn_auto_machine_weight: float = 1.0,185        style_fidelity: float = 0.5,186        reference_attn: bool = True,187        reference_adain: bool = True,188    ):189        r"""190        Function invoked when calling the pipeline for generation.191 192        Args:193            prompt (`str` or `List[str]`, *optional*):194                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.195                instead.196            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):197                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If198                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can199                also be accepted as an image.200            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):201                The height in pixels of the generated image.202            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):203                The width in pixels of the generated image.204            num_inference_steps (`int`, *optional*, defaults to 50):205                The number of denoising steps. More denoising steps usually lead to a higher quality image at the206                expense of slower inference.207            guidance_scale (`float`, *optional*, defaults to 7.5):208                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).209                `guidance_scale` is defined as `w` of equation 2. of [Imagen210                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >211                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,212                usually at the expense of lower image quality.213            negative_prompt (`str` or `List[str]`, *optional*):214                The prompt or prompts not to guide the image generation. If not defined, one has to pass215                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is216                less than `1`).217            num_images_per_prompt (`int`, *optional*, defaults to 1):218                The number of images to generate per prompt.219            eta (`float`, *optional*, defaults to 0.0):220                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to221                [`schedulers.DDIMScheduler`], will be ignored for others.222            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):223                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)224                to make generation deterministic.225            latents (`torch.FloatTensor`, *optional*):226                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image227                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents228                tensor will ge generated by sampling using the supplied random `generator`.229            prompt_embeds (`torch.FloatTensor`, *optional*):230                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not231                provided, text embeddings will be generated from `prompt` input argument.232            negative_prompt_embeds (`torch.FloatTensor`, *optional*):233                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt234                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input235                argument.236            output_type (`str`, *optional*, defaults to `"pil"`):237                The output format of the generate image. Choose between238                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.239            return_dict (`bool`, *optional*, defaults to `True`):240                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a241                plain tuple.242            callback (`Callable`, *optional*):243                A function that will be called every `callback_steps` steps during inference. The function will be244                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.245            callback_steps (`int`, *optional*, defaults to 1):246                The frequency at which the `callback` function will be called. If not specified, the callback will be247                called at every step.248            cross_attention_kwargs (`dict`, *optional*):249                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under250                `self.processor` in251                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).252            guidance_rescale (`float`, *optional*, defaults to 0.0):253                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are254                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of255                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).256                Guidance rescale factor should fix overexposure when using zero terminal SNR.257            attention_auto_machine_weight (`float`):258                Weight of using reference query for self attention's context.259                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.260            gn_auto_machine_weight (`float`):261                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.262            style_fidelity (`float`):263                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,264                elif style_fidelity=0.0, prompt more important, else balanced.265            reference_attn (`bool`):266                Whether to use reference query for self attention's context.267            reference_adain (`bool`):268                Whether to use reference adain.269 270        Examples:271 272        Returns:273            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:274            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.275            When returning a tuple, the first element is a list with the generated images, and the second element is a276            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"277            (nsfw) content, according to the `safety_checker`.278        """279        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."280 281        # 0. Default height and width to unet282        height, width = self._default_height_width(height, width, ref_image)283 284        # 1. Check inputs. Raise error if not correct285        self.check_inputs(286            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds287        )288 289        # 2. Define call parameters290        if prompt is not None and isinstance(prompt, str):291            batch_size = 1292        elif prompt is not None and isinstance(prompt, list):293            batch_size = len(prompt)294        else:295            batch_size = prompt_embeds.shape[0]296 297        device = self._execution_device298        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)299        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`300        # corresponds to doing no classifier free guidance.301        do_classifier_free_guidance = guidance_scale > 1.0302 303        # 3. Encode input prompt304        text_encoder_lora_scale = (305            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None306        )307        prompt_embeds = self._encode_prompt(308            prompt,309            device,310            num_images_per_prompt,311            do_classifier_free_guidance,312            negative_prompt,313            prompt_embeds=prompt_embeds,314            negative_prompt_embeds=negative_prompt_embeds,315            lora_scale=text_encoder_lora_scale,316        )317 318        # 4. Preprocess reference image319        ref_image = self.prepare_image(320            image=ref_image,321            width=width,322            height=height,323            batch_size=batch_size * num_images_per_prompt,324            num_images_per_prompt=num_images_per_prompt,325            device=device,326            dtype=prompt_embeds.dtype,327        )328 329        # 5. Prepare timesteps330        self.scheduler.set_timesteps(num_inference_steps, device=device)331        timesteps = self.scheduler.timesteps332 333        # 6. Prepare latent variables334        num_channels_latents = self.unet.config.in_channels335        latents = self.prepare_latents(336            batch_size * num_images_per_prompt,337            num_channels_latents,338            height,339            width,340            prompt_embeds.dtype,341            device,342            generator,343            latents,344        )345 346        # 7. Prepare reference latent variables347        ref_image_latents = self.prepare_ref_latents(348            ref_image,349            batch_size * num_images_per_prompt,350            prompt_embeds.dtype,351            device,352            generator,353            do_classifier_free_guidance,354        )355 356        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline357        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)358 359        # 9. Modify self attention and group norm360        MODE = "write"361        uc_mask = (362            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)363            .type_as(ref_image_latents)364            .bool()365        )366 367        def hacked_basic_transformer_inner_forward(368            self,369            hidden_states: torch.FloatTensor,370            attention_mask: Optional[torch.FloatTensor] = None,371            encoder_hidden_states: Optional[torch.FloatTensor] = None,372            encoder_attention_mask: Optional[torch.FloatTensor] = None,373            timestep: Optional[torch.LongTensor] = None,374            cross_attention_kwargs: Dict[str, Any] = None,375            class_labels: Optional[torch.LongTensor] = None,376        ):377            if self.use_ada_layer_norm:378                norm_hidden_states = self.norm1(hidden_states, timestep)379            elif self.use_ada_layer_norm_zero:380                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(381                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype382                )383            else:384                norm_hidden_states = self.norm1(hidden_states)385 386            # 1. Self-Attention387            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}388            if self.only_cross_attention:389                attn_output = self.attn1(390                    norm_hidden_states,391                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,392                    attention_mask=attention_mask,393                    **cross_attention_kwargs,394                )395            else:396                if MODE == "write":397                    self.bank.append(norm_hidden_states.detach().clone())398                    attn_output = self.attn1(399                        norm_hidden_states,400                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,401                        attention_mask=attention_mask,402                        **cross_attention_kwargs,403                    )404                if MODE == "read":405                    if attention_auto_machine_weight > self.attn_weight:406                        attn_output_uc = self.attn1(407                            norm_hidden_states,408                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),409                            # attention_mask=attention_mask,410                            **cross_attention_kwargs,411                        )412                        attn_output_c = attn_output_uc.clone()413                        if do_classifier_free_guidance and style_fidelity > 0:414                            attn_output_c[uc_mask] = self.attn1(415                                norm_hidden_states[uc_mask],416                                encoder_hidden_states=norm_hidden_states[uc_mask],417                                **cross_attention_kwargs,418                            )419                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc420                        self.bank.clear()421                    else:422                        attn_output = self.attn1(423                            norm_hidden_states,424                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,425                            attention_mask=attention_mask,426                            **cross_attention_kwargs,427                        )428            if self.use_ada_layer_norm_zero:429                attn_output = gate_msa.unsqueeze(1) * attn_output430            hidden_states = attn_output + hidden_states431 432            if self.attn2 is not None:433                norm_hidden_states = (434                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)435                )436 437                # 2. Cross-Attention438                attn_output = self.attn2(439                    norm_hidden_states,440                    encoder_hidden_states=encoder_hidden_states,441                    attention_mask=encoder_attention_mask,442                    **cross_attention_kwargs,443                )444                hidden_states = attn_output + hidden_states445 446            # 3. Feed-forward447            norm_hidden_states = self.norm3(hidden_states)448 449            if self.use_ada_layer_norm_zero:450                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]451 452            ff_output = self.ff(norm_hidden_states)453 454            if self.use_ada_layer_norm_zero:455                ff_output = gate_mlp.unsqueeze(1) * ff_output456 457            hidden_states = ff_output + hidden_states458 459            return hidden_states460 461        def hacked_mid_forward(self, *args, **kwargs):462            eps = 1e-6463            x = self.original_forward(*args, **kwargs)464            if MODE == "write":465                if gn_auto_machine_weight >= self.gn_weight:466                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)467                    self.mean_bank.append(mean)468                    self.var_bank.append(var)469            if MODE == "read":470                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:471                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)472                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5473                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))474                    var_acc = sum(self.var_bank) / float(len(self.var_bank))475                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5476                    x_uc = (((x - mean) / std) * std_acc) + mean_acc477                    x_c = x_uc.clone()478                    if do_classifier_free_guidance and style_fidelity > 0:479                        x_c[uc_mask] = x[uc_mask]480                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc481                self.mean_bank = []482                self.var_bank = []483            return x484 485        def hack_CrossAttnDownBlock2D_forward(486            self,487            hidden_states: torch.FloatTensor,488            temb: Optional[torch.FloatTensor] = None,489            encoder_hidden_states: Optional[torch.FloatTensor] = None,490            attention_mask: Optional[torch.FloatTensor] = None,491            cross_attention_kwargs: Optional[Dict[str, Any]] = None,492            encoder_attention_mask: Optional[torch.FloatTensor] = None,493        ):494            eps = 1e-6495 496            # TODO(Patrick, William) - attention mask is not used497            output_states = ()498 499            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):500                hidden_states = resnet(hidden_states, temb)501                hidden_states = attn(502                    hidden_states,503                    encoder_hidden_states=encoder_hidden_states,504                    cross_attention_kwargs=cross_attention_kwargs,505                    attention_mask=attention_mask,506                    encoder_attention_mask=encoder_attention_mask,507                    return_dict=False,508                )[0]509                if MODE == "write":510                    if gn_auto_machine_weight >= self.gn_weight:511                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)512                        self.mean_bank.append([mean])513                        self.var_bank.append([var])514                if MODE == "read":515                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:516                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)517                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5518                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))519                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))520                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5521                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc522                        hidden_states_c = hidden_states_uc.clone()523                        if do_classifier_free_guidance and style_fidelity > 0:524                            hidden_states_c[uc_mask] = hidden_states[uc_mask]525                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc526 527                output_states = output_states + (hidden_states,)528 529            if MODE == "read":530                self.mean_bank = []531                self.var_bank = []532 533            if self.downsamplers is not None:534                for downsampler in self.downsamplers:535                    hidden_states = downsampler(hidden_states)536 537                output_states = output_states + (hidden_states,)538 539            return hidden_states, output_states540 541        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):542            eps = 1e-6543 544            output_states = ()545 546            for i, resnet in enumerate(self.resnets):547                hidden_states = resnet(hidden_states, temb)548 549                if MODE == "write":550                    if gn_auto_machine_weight >= self.gn_weight:551                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)552                        self.mean_bank.append([mean])553                        self.var_bank.append([var])554                if MODE == "read":555                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:556                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)557                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5558                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))559                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))560                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5561                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc562                        hidden_states_c = hidden_states_uc.clone()563                        if do_classifier_free_guidance and style_fidelity > 0:564                            hidden_states_c[uc_mask] = hidden_states[uc_mask]565                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc566 567                output_states = output_states + (hidden_states,)568 569            if MODE == "read":570                self.mean_bank = []571                self.var_bank = []572 573            if self.downsamplers is not None:574                for downsampler in self.downsamplers:575                    hidden_states = downsampler(hidden_states)576 577                output_states = output_states + (hidden_states,)578 579            return hidden_states, output_states580 581        def hacked_CrossAttnUpBlock2D_forward(582            self,583            hidden_states: torch.FloatTensor,584            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],585            temb: Optional[torch.FloatTensor] = None,586            encoder_hidden_states: Optional[torch.FloatTensor] = None,587            cross_attention_kwargs: Optional[Dict[str, Any]] = None,588            upsample_size: Optional[int] = None,589            attention_mask: Optional[torch.FloatTensor] = None,590            encoder_attention_mask: Optional[torch.FloatTensor] = None,591        ):592            eps = 1e-6593            # TODO(Patrick, William) - attention mask is not used594            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):595                # pop res hidden states596                res_hidden_states = res_hidden_states_tuple[-1]597                res_hidden_states_tuple = res_hidden_states_tuple[:-1]598                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)599                hidden_states = resnet(hidden_states, temb)600                hidden_states = attn(601                    hidden_states,602                    encoder_hidden_states=encoder_hidden_states,603                    cross_attention_kwargs=cross_attention_kwargs,604                    attention_mask=attention_mask,605                    encoder_attention_mask=encoder_attention_mask,606                    return_dict=False,607                )[0]608 609                if MODE == "write":610                    if gn_auto_machine_weight >= self.gn_weight:611                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)612                        self.mean_bank.append([mean])613                        self.var_bank.append([var])614                if MODE == "read":615                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:616                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)617                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5618                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))619                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))620                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5621                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc622                        hidden_states_c = hidden_states_uc.clone()623                        if do_classifier_free_guidance and style_fidelity > 0:624                            hidden_states_c[uc_mask] = hidden_states[uc_mask]625                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc626 627            if MODE == "read":628                self.mean_bank = []629                self.var_bank = []630 631            if self.upsamplers is not None:632                for upsampler in self.upsamplers:633                    hidden_states = upsampler(hidden_states, upsample_size)634 635            return hidden_states636 637        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):638            eps = 1e-6639            for i, resnet in enumerate(self.resnets):640                # pop res hidden states641                res_hidden_states = res_hidden_states_tuple[-1]642                res_hidden_states_tuple = res_hidden_states_tuple[:-1]643                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)644                hidden_states = resnet(hidden_states, temb)645 646                if MODE == "write":647                    if gn_auto_machine_weight >= self.gn_weight:648                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)649                        self.mean_bank.append([mean])650                        self.var_bank.append([var])651                if MODE == "read":652                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:653                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)654                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5655                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))656                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))657                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5658                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc659                        hidden_states_c = hidden_states_uc.clone()660                        if do_classifier_free_guidance and style_fidelity > 0:661                            hidden_states_c[uc_mask] = hidden_states[uc_mask]662                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc663 664            if MODE == "read":665                self.mean_bank = []666                self.var_bank = []667 668            if self.upsamplers is not None:669                for upsampler in self.upsamplers:670                    hidden_states = upsampler(hidden_states, upsample_size)671 672            return hidden_states673 674        if reference_attn:675            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]676            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])677 678            for i, module in enumerate(attn_modules):679                module._original_inner_forward = module.forward680                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)681                module.bank = []682                module.attn_weight = float(i) / float(len(attn_modules))683 684        if reference_adain:685            gn_modules = [self.unet.mid_block]686            self.unet.mid_block.gn_weight = 0687 688            down_blocks = self.unet.down_blocks689            for w, module in enumerate(down_blocks):690                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))691                gn_modules.append(module)692 693            up_blocks = self.unet.up_blocks694            for w, module in enumerate(up_blocks):695                module.gn_weight = float(w) / float(len(up_blocks))696                gn_modules.append(module)697 698            for i, module in enumerate(gn_modules):699                if getattr(module, "original_forward", None) is None:700                    module.original_forward = module.forward701                if i == 0:702                    # mid_block703                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)704                elif isinstance(module, CrossAttnDownBlock2D):705                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)706                elif isinstance(module, DownBlock2D):707                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)708                elif isinstance(module, CrossAttnUpBlock2D):709                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)710                elif isinstance(module, UpBlock2D):711                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)712                module.mean_bank = []713                module.var_bank = []714                module.gn_weight *= 2715 716        # 10. Denoising loop717        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order718        with self.progress_bar(total=num_inference_steps) as progress_bar:719            for i, t in enumerate(timesteps):720                # expand the latents if we are doing classifier free guidance721                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents722                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)723 724                # ref only part725                noise = randn_tensor(726                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype727                )728                ref_xt = self.scheduler.add_noise(729                    ref_image_latents,730                    noise,731                    t.reshape(732                        1,733                    ),734                )735                ref_xt = torch.cat([ref_xt] * 2) if do_classifier_free_guidance else ref_xt736                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                        step_idx = i // getattr(self.scheduler, "order", 1)774                        callback(step_idx, t, latents)775 776        if not output_type == "latent":777            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]778            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)779        else:780            image = latents781            has_nsfw_concept = None782 783        if has_nsfw_concept is None:784            do_denormalize = [True] * image.shape[0]785        else:786            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]787 788        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)789 790        # Offload last model to CPU791        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:792            self.final_offload_hook.offload()793 794        if not return_dict:795            return (image, has_nsfw_concept)796 797        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)798