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.utils import PIL_INTERPOLATION, logging, randn_tensor13 14 15logger = logging.get_logger(__name__)  # pylint: disable=invalid-name16 17EXAMPLE_DOC_STRING = """18    Examples:19        ```py20        >>> import torch21        >>> from diffusers import UniPCMultistepScheduler22        >>> from diffusers.utils import load_image23 24        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")25 26        >>> pipe = StableDiffusionReferencePipeline.from_pretrained(27                "runwayml/stable-diffusion-v1-5",28                safety_checker=None,29                torch_dtype=torch.float1630                ).to('cuda:0')31 32        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)33 34        >>> result_img = pipe(ref_image=input_image,35                        prompt="1girl",36                        num_inference_steps=20,37                        reference_attn=True,38                        reference_adain=True).images[0]39 40        >>> result_img.show()41        ```42"""43 44 45def torch_dfs(model: torch.nn.Module):46    result = [model]47    for child in model.children():48        result += torch_dfs(child)49    return result50 51 52class StableDiffusionReferencePipeline(StableDiffusionPipeline):53    def _default_height_width(self, height, width, image):54        # NOTE: It is possible that a list of images have different55        # dimensions for each image, so just checking the first image56        # is not _exactly_ correct, but it is simple.57        while isinstance(image, list):58            image = image[0]59 60        if height is None:61            if isinstance(image, PIL.Image.Image):62                height = image.height63            elif isinstance(image, torch.Tensor):64                height = image.shape[2]65 66            height = (height // 8) * 8  # round down to nearest multiple of 867 68        if width is None:69            if isinstance(image, PIL.Image.Image):70                width = image.width71            elif isinstance(image, torch.Tensor):72                width = image.shape[3]73 74            width = (width // 8) * 8  # round down to nearest multiple of 875 76        return height, width77 78    def prepare_image(79        self,80        image,81        width,82        height,83        batch_size,84        num_images_per_prompt,85        device,86        dtype,87        do_classifier_free_guidance=False,88        guess_mode=False,89    ):90        if not isinstance(image, torch.Tensor):91            if isinstance(image, PIL.Image.Image):92                image = [image]93 94            if isinstance(image[0], PIL.Image.Image):95                images = []96 97                for image_ in image:98                    image_ = image_.convert("RGB")99                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])100                    image_ = np.array(image_)101                    image_ = image_[None, :]102                    images.append(image_)103 104                image = images105 106                image = np.concatenate(image, axis=0)107                image = np.array(image).astype(np.float32) / 255.0108                image = (image - 0.5) / 0.5109                image = image.transpose(0, 3, 1, 2)110                image = torch.from_numpy(image)111            elif isinstance(image[0], torch.Tensor):112                image = torch.cat(image, dim=0)113 114        image_batch_size = image.shape[0]115 116        if image_batch_size == 1:117            repeat_by = batch_size118        else:119            # image batch size is the same as prompt batch size120            repeat_by = num_images_per_prompt121 122        image = image.repeat_interleave(repeat_by, dim=0)123 124        image = image.to(device=device, dtype=dtype)125 126        if do_classifier_free_guidance and not guess_mode:127            image = torch.cat([image] * 2)128 129        return image130 131    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):132        refimage = refimage.to(device=device, dtype=dtype)133 134        # encode the mask image into latents space so we can concatenate it to the latents135        if isinstance(generator, list):136            ref_image_latents = [137                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])138                for i in range(batch_size)139            ]140            ref_image_latents = torch.cat(ref_image_latents, dim=0)141        else:142            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)143        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents144 145        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method146        if ref_image_latents.shape[0] < batch_size:147            if not batch_size % ref_image_latents.shape[0] == 0:148                raise ValueError(149                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"150                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."151                    " Make sure the number of images that you pass is divisible by the total requested batch size."152                )153            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)154 155        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents156 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        attention_auto_machine_weight: float = 1.0,183        gn_auto_machine_weight: float = 1.0,184        style_fidelity: float = 0.5,185        reference_attn: bool = True,186        reference_adain: bool = True,187    ):188        r"""189        Function invoked when calling the pipeline for generation.190 191        Args:192            prompt (`str` or `List[str]`, *optional*):193                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.194                instead.195            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):196                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If197                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can198                also be accepted as an image.199            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):200                The height in pixels of the generated image.201            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):202                The width in pixels of the generated image.203            num_inference_steps (`int`, *optional*, defaults to 50):204                The number of denoising steps. More denoising steps usually lead to a higher quality image at the205                expense of slower inference.206            guidance_scale (`float`, *optional*, defaults to 7.5):207                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).208                `guidance_scale` is defined as `w` of equation 2. of [Imagen209                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >210                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,211                usually at the expense of lower image quality.212            negative_prompt (`str` or `List[str]`, *optional*):213                The prompt or prompts not to guide the image generation. If not defined, one has to pass214                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is215                less than `1`).216            num_images_per_prompt (`int`, *optional*, defaults to 1):217                The number of images to generate per prompt.218            eta (`float`, *optional*, defaults to 0.0):219                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to220                [`schedulers.DDIMScheduler`], will be ignored for others.221            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):222                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)223                to make generation deterministic.224            latents (`torch.FloatTensor`, *optional*):225                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image226                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents227                tensor will ge generated by sampling using the supplied random `generator`.228            prompt_embeds (`torch.FloatTensor`, *optional*):229                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not230                provided, text embeddings will be generated from `prompt` input argument.231            negative_prompt_embeds (`torch.FloatTensor`, *optional*):232                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt233                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input234                argument.235            output_type (`str`, *optional*, defaults to `"pil"`):236                The output format of the generate image. Choose between237                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.238            return_dict (`bool`, *optional*, defaults to `True`):239                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a240                plain tuple.241            callback (`Callable`, *optional*):242                A function that will be called every `callback_steps` steps during inference. The function will be243                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.244            callback_steps (`int`, *optional*, defaults to 1):245                The frequency at which the `callback` function will be called. If not specified, the callback will be246                called at every step.247            cross_attention_kwargs (`dict`, *optional*):248                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under249                `self.processor` in250                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).251            attention_auto_machine_weight (`float`):252                Weight of using reference query for self attention's context.253                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.254            gn_auto_machine_weight (`float`):255                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.256            style_fidelity (`float`):257                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,258                elif style_fidelity=0.0, prompt more important, else balanced.259            reference_attn (`bool`):260                Whether to use reference query for self attention's context.261            reference_adain (`bool`):262                Whether to use reference adain.263 264        Examples:265 266        Returns:267            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:268            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.269            When returning a tuple, the first element is a list with the generated images, and the second element is a270            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"271            (nsfw) content, according to the `safety_checker`.272        """273        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."274 275        # 0. Default height and width to unet276        height, width = self._default_height_width(height, width, ref_image)277 278        # 1. Check inputs. Raise error if not correct279        self.check_inputs(280            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds281        )282 283        # 2. Define call parameters284        if prompt is not None and isinstance(prompt, str):285            batch_size = 1286        elif prompt is not None and isinstance(prompt, list):287            batch_size = len(prompt)288        else:289            batch_size = prompt_embeds.shape[0]290 291        device = self._execution_device292        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)293        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`294        # corresponds to doing no classifier free guidance.295        do_classifier_free_guidance = guidance_scale > 1.0296 297        # 3. Encode input prompt298        prompt_embeds = self._encode_prompt(299            prompt,300            device,301            num_images_per_prompt,302            do_classifier_free_guidance,303            negative_prompt,304            prompt_embeds=prompt_embeds,305            negative_prompt_embeds=negative_prompt_embeds,306        )307 308        # 4. Preprocess reference image309        ref_image = self.prepare_image(310            image=ref_image,311            width=width,312            height=height,313            batch_size=batch_size * num_images_per_prompt,314            num_images_per_prompt=num_images_per_prompt,315            device=device,316            dtype=prompt_embeds.dtype,317        )318 319        # 5. Prepare timesteps320        self.scheduler.set_timesteps(num_inference_steps, device=device)321        timesteps = self.scheduler.timesteps322 323        # 6. Prepare latent variables324        num_channels_latents = self.unet.config.in_channels325        latents = self.prepare_latents(326            batch_size * num_images_per_prompt,327            num_channels_latents,328            height,329            width,330            prompt_embeds.dtype,331            device,332            generator,333            latents,334        )335 336        # 7. Prepare reference latent variables337        ref_image_latents = self.prepare_ref_latents(338            ref_image,339            batch_size * num_images_per_prompt,340            prompt_embeds.dtype,341            device,342            generator,343            do_classifier_free_guidance,344        )345 346        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline347        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)348 349        # 9. Modify self attention and group norm350        MODE = "write"351        uc_mask = (352            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)353            .type_as(ref_image_latents)354            .bool()355        )356 357        def hacked_basic_transformer_inner_forward(358            self,359            hidden_states: torch.FloatTensor,360            attention_mask: Optional[torch.FloatTensor] = None,361            encoder_hidden_states: Optional[torch.FloatTensor] = None,362            encoder_attention_mask: Optional[torch.FloatTensor] = None,363            timestep: Optional[torch.LongTensor] = None,364            cross_attention_kwargs: Dict[str, Any] = None,365            class_labels: Optional[torch.LongTensor] = None,366        ):367            if self.use_ada_layer_norm:368                norm_hidden_states = self.norm1(hidden_states, timestep)369            elif self.use_ada_layer_norm_zero:370                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(371                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype372                )373            else:374                norm_hidden_states = self.norm1(hidden_states)375 376            # 1. Self-Attention377            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}378            if self.only_cross_attention:379                attn_output = self.attn1(380                    norm_hidden_states,381                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,382                    attention_mask=attention_mask,383                    **cross_attention_kwargs,384                )385            else:386                if MODE == "write":387                    self.bank.append(norm_hidden_states.detach().clone())388                    attn_output = self.attn1(389                        norm_hidden_states,390                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,391                        attention_mask=attention_mask,392                        **cross_attention_kwargs,393                    )394                if MODE == "read":395                    if attention_auto_machine_weight > self.attn_weight:396                        attn_output_uc = self.attn1(397                            norm_hidden_states,398                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),399                            # attention_mask=attention_mask,400                            **cross_attention_kwargs,401                        )402                        attn_output_c = attn_output_uc.clone()403                        if do_classifier_free_guidance and style_fidelity > 0:404                            attn_output_c[uc_mask] = self.attn1(405                                norm_hidden_states[uc_mask],406                                encoder_hidden_states=norm_hidden_states[uc_mask],407                                **cross_attention_kwargs,408                            )409                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc410                        self.bank.clear()411                    else:412                        attn_output = self.attn1(413                            norm_hidden_states,414                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,415                            attention_mask=attention_mask,416                            **cross_attention_kwargs,417                        )418            if self.use_ada_layer_norm_zero:419                attn_output = gate_msa.unsqueeze(1) * attn_output420            hidden_states = attn_output + hidden_states421 422            if self.attn2 is not None:423                norm_hidden_states = (424                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)425                )426 427                # 2. Cross-Attention428                attn_output = self.attn2(429                    norm_hidden_states,430                    encoder_hidden_states=encoder_hidden_states,431                    attention_mask=encoder_attention_mask,432                    **cross_attention_kwargs,433                )434                hidden_states = attn_output + hidden_states435 436            # 3. Feed-forward437            norm_hidden_states = self.norm3(hidden_states)438 439            if self.use_ada_layer_norm_zero:440                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]441 442            ff_output = self.ff(norm_hidden_states)443 444            if self.use_ada_layer_norm_zero:445                ff_output = gate_mlp.unsqueeze(1) * ff_output446 447            hidden_states = ff_output + hidden_states448 449            return hidden_states450 451        def hacked_mid_forward(self, *args, **kwargs):452            eps = 1e-6453            x = self.original_forward(*args, **kwargs)454            if MODE == "write":455                if gn_auto_machine_weight >= self.gn_weight:456                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)457                    self.mean_bank.append(mean)458                    self.var_bank.append(var)459            if MODE == "read":460                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:461                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)462                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5463                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))464                    var_acc = sum(self.var_bank) / float(len(self.var_bank))465                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5466                    x_uc = (((x - mean) / std) * std_acc) + mean_acc467                    x_c = x_uc.clone()468                    if do_classifier_free_guidance and style_fidelity > 0:469                        x_c[uc_mask] = x[uc_mask]470                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc471                self.mean_bank = []472                self.var_bank = []473            return x474 475        def hack_CrossAttnDownBlock2D_forward(476            self,477            hidden_states: torch.FloatTensor,478            temb: Optional[torch.FloatTensor] = None,479            encoder_hidden_states: Optional[torch.FloatTensor] = None,480            attention_mask: Optional[torch.FloatTensor] = None,481            cross_attention_kwargs: Optional[Dict[str, Any]] = None,482            encoder_attention_mask: Optional[torch.FloatTensor] = None,483        ):484            eps = 1e-6485 486            # TODO(Patrick, William) - attention mask is not used487            output_states = ()488 489            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):490                hidden_states = resnet(hidden_states, temb)491                hidden_states = attn(492                    hidden_states,493                    encoder_hidden_states=encoder_hidden_states,494                    cross_attention_kwargs=cross_attention_kwargs,495                    attention_mask=attention_mask,496                    encoder_attention_mask=encoder_attention_mask,497                    return_dict=False,498                )[0]499                if MODE == "write":500                    if gn_auto_machine_weight >= self.gn_weight:501                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)502                        self.mean_bank.append([mean])503                        self.var_bank.append([var])504                if MODE == "read":505                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:506                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)507                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5508                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))509                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))510                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5511                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc512                        hidden_states_c = hidden_states_uc.clone()513                        if do_classifier_free_guidance and style_fidelity > 0:514                            hidden_states_c[uc_mask] = hidden_states[uc_mask]515                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc516 517                output_states = output_states + (hidden_states,)518 519            if MODE == "read":520                self.mean_bank = []521                self.var_bank = []522 523            if self.downsamplers is not None:524                for downsampler in self.downsamplers:525                    hidden_states = downsampler(hidden_states)526 527                output_states = output_states + (hidden_states,)528 529            return hidden_states, output_states530 531        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):532            eps = 1e-6533 534            output_states = ()535 536            for i, resnet in enumerate(self.resnets):537                hidden_states = resnet(hidden_states, temb)538 539                if MODE == "write":540                    if gn_auto_machine_weight >= self.gn_weight:541                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)542                        self.mean_bank.append([mean])543                        self.var_bank.append([var])544                if MODE == "read":545                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:546                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)547                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5548                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))549                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))550                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5551                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc552                        hidden_states_c = hidden_states_uc.clone()553                        if do_classifier_free_guidance and style_fidelity > 0:554                            hidden_states_c[uc_mask] = hidden_states[uc_mask]555                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc556 557                output_states = output_states + (hidden_states,)558 559            if MODE == "read":560                self.mean_bank = []561                self.var_bank = []562 563            if self.downsamplers is not None:564                for downsampler in self.downsamplers:565                    hidden_states = downsampler(hidden_states)566 567                output_states = output_states + (hidden_states,)568 569            return hidden_states, output_states570 571        def hacked_CrossAttnUpBlock2D_forward(572            self,573            hidden_states: torch.FloatTensor,574            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],575            temb: Optional[torch.FloatTensor] = None,576            encoder_hidden_states: Optional[torch.FloatTensor] = None,577            cross_attention_kwargs: Optional[Dict[str, Any]] = None,578            upsample_size: Optional[int] = None,579            attention_mask: Optional[torch.FloatTensor] = None,580            encoder_attention_mask: Optional[torch.FloatTensor] = None,581        ):582            eps = 1e-6583            # TODO(Patrick, William) - attention mask is not used584            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):585                # pop res hidden states586                res_hidden_states = res_hidden_states_tuple[-1]587                res_hidden_states_tuple = res_hidden_states_tuple[:-1]588                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)589                hidden_states = resnet(hidden_states, temb)590                hidden_states = attn(591                    hidden_states,592                    encoder_hidden_states=encoder_hidden_states,593                    cross_attention_kwargs=cross_attention_kwargs,594                    attention_mask=attention_mask,595                    encoder_attention_mask=encoder_attention_mask,596                    return_dict=False,597                )[0]598 599                if MODE == "write":600                    if gn_auto_machine_weight >= self.gn_weight:601                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)602                        self.mean_bank.append([mean])603                        self.var_bank.append([var])604                if MODE == "read":605                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:606                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)607                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5608                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))609                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))610                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5611                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc612                        hidden_states_c = hidden_states_uc.clone()613                        if do_classifier_free_guidance and style_fidelity > 0:614                            hidden_states_c[uc_mask] = hidden_states[uc_mask]615                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc616 617            if MODE == "read":618                self.mean_bank = []619                self.var_bank = []620 621            if self.upsamplers is not None:622                for upsampler in self.upsamplers:623                    hidden_states = upsampler(hidden_states, upsample_size)624 625            return hidden_states626 627        def hacked_UpBlock2D_forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):628            eps = 1e-6629            for i, resnet in enumerate(self.resnets):630                # pop res hidden states631                res_hidden_states = res_hidden_states_tuple[-1]632                res_hidden_states_tuple = res_hidden_states_tuple[:-1]633                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)634                hidden_states = resnet(hidden_states, temb)635 636                if MODE == "write":637                    if gn_auto_machine_weight >= self.gn_weight:638                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)639                        self.mean_bank.append([mean])640                        self.var_bank.append([var])641                if MODE == "read":642                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:643                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)644                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5645                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))646                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))647                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5648                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc649                        hidden_states_c = hidden_states_uc.clone()650                        if do_classifier_free_guidance and style_fidelity > 0:651                            hidden_states_c[uc_mask] = hidden_states[uc_mask]652                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc653 654            if MODE == "read":655                self.mean_bank = []656                self.var_bank = []657 658            if self.upsamplers is not None:659                for upsampler in self.upsamplers:660                    hidden_states = upsampler(hidden_states, upsample_size)661 662            return hidden_states663 664        if reference_attn:665            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]666            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])667 668            for i, module in enumerate(attn_modules):669                module._original_inner_forward = module.forward670                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)671                module.bank = []672                module.attn_weight = float(i) / float(len(attn_modules))673 674        if reference_adain:675            gn_modules = [self.unet.mid_block]676            self.unet.mid_block.gn_weight = 0677 678            down_blocks = self.unet.down_blocks679            for w, module in enumerate(down_blocks):680                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))681                gn_modules.append(module)682 683            up_blocks = self.unet.up_blocks684            for w, module in enumerate(up_blocks):685                module.gn_weight = float(w) / float(len(up_blocks))686                gn_modules.append(module)687 688            for i, module in enumerate(gn_modules):689                if getattr(module, "original_forward", None) is None:690                    module.original_forward = module.forward691                if i == 0:692                    # mid_block693                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)694                elif isinstance(module, CrossAttnDownBlock2D):695                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)696                elif isinstance(module, DownBlock2D):697                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)698                elif isinstance(module, CrossAttnUpBlock2D):699                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)700                elif isinstance(module, UpBlock2D):701                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)702                module.mean_bank = []703                module.var_bank = []704                module.gn_weight *= 2705 706        # 10. Denoising loop707        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order708        with self.progress_bar(total=num_inference_steps) as progress_bar:709            for i, t in enumerate(timesteps):710                # expand the latents if we are doing classifier free guidance711                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents712                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)713 714                # ref only part715                noise = randn_tensor(716                    ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype717                )718                ref_xt = self.scheduler.add_noise(719                    ref_image_latents,720                    noise,721                    t.reshape(722                        1,723                    ),724                )725                ref_xt = self.scheduler.scale_model_input(ref_xt, t)726 727                MODE = "write"728                self.unet(729                    ref_xt,730                    t,731                    encoder_hidden_states=prompt_embeds,732                    cross_attention_kwargs=cross_attention_kwargs,733                    return_dict=False,734                )735 736                # predict the noise residual737                MODE = "read"738                noise_pred = self.unet(739                    latent_model_input,740                    t,741                    encoder_hidden_states=prompt_embeds,742                    cross_attention_kwargs=cross_attention_kwargs,743                    return_dict=False,744                )[0]745 746                # perform guidance747                if do_classifier_free_guidance:748                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)749                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)750 751                # compute the previous noisy sample x_t -> x_t-1752                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]753 754                # call the callback, if provided755                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):756                    progress_bar.update()757                    if callback is not None and i % callback_steps == 0:758                        callback(i, t, latents)759 760        if not output_type == "latent":761            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]762            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)763        else:764            image = latents765            has_nsfw_concept = None766 767        if has_nsfw_concept is None:768            do_denormalize = [True] * image.shape[0]769        else:770            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]771 772        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)773 774        # Offload last model to CPU775        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:776            self.final_offload_hook.offload()777 778        if not return_dict:779            return (image, has_nsfw_concept)780 781        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)782