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/12802import inspect3from typing import Any, Callable, Dict, List, Optional, Tuple, Union4 5import numpy as np6import PIL.Image7import torch8from packaging import version9from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer10 11from diffusers import AutoencoderKL, DiffusionPipeline, UNet2DConditionModel12from diffusers.configuration_utils import FrozenDict, deprecate13from diffusers.image_processor import VaeImageProcessor14from diffusers.loaders import (15    FromSingleFileMixin,16    IPAdapterMixin,17    StableDiffusionLoraLoaderMixin,18    TextualInversionLoaderMixin,19)20from diffusers.models.attention import BasicTransformerBlock21from diffusers.models.lora import adjust_lora_scale_text_encoder22from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D23from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput24from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg25from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker26from diffusers.schedulers import KarrasDiffusionSchedulers27from diffusers.utils import (28    PIL_INTERPOLATION,29    USE_PEFT_BACKEND,30    logging,31    scale_lora_layers,32    unscale_lora_layers,33)34from diffusers.utils.torch_utils import randn_tensor35 36 37logger = logging.get_logger(__name__)  # pylint: disable=invalid-name38 39EXAMPLE_DOC_STRING = """40    Examples:41        ```py42        >>> import torch43        >>> from diffusers import UniPCMultistepScheduler44        >>> from diffusers.utils import load_image45 46        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")47 48        >>> pipe = StableDiffusionReferencePipeline.from_pretrained(49                "runwayml/stable-diffusion-v1-5",50                safety_checker=None,51                torch_dtype=torch.float1652                ).to('cuda:0')53 54        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)55 56        >>> result_img = pipe(ref_image=input_image,57                        prompt="1girl",58                        num_inference_steps=20,59                        reference_attn=True,60                        reference_adain=True).images[0]61 62        >>> result_img.show()63        ```64"""65 66 67def torch_dfs(model: torch.nn.Module):68    r"""69    Performs a depth-first search on the given PyTorch model and returns a list of all its child modules.70 71    Args:72        model (torch.nn.Module): The PyTorch model to perform the depth-first search on.73 74    Returns:75        list: A list of all child modules of the given model.76    """77    result = [model]78    for child in model.children():79        result += torch_dfs(child)80    return result81 82 83class StableDiffusionReferencePipeline(84    DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin85):86    r"""87    Pipeline for Stable Diffusion Reference.88 89    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods90    implemented for all pipelines (downloading, saving, running on a particular device, etc.).91 92    The pipeline also inherits the following loading methods:93    - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings94    - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights95    - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights96    - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files97    - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters98 99    Args:100        vae ([`AutoencoderKL`]):101            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.102        text_encoder ([`CLIPTextModel`]):103            Frozen text-encoder. Stable Diffusion uses the text portion of104            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically105            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.106        tokenizer (`CLIPTokenizer`):107            Tokenizer of class108            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).109        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.110        scheduler ([`SchedulerMixin`]):111            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of112            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].113        safety_checker ([`StableDiffusionSafetyChecker`]):114            Classification module that estimates whether generated images could be considered offensive or harmful.115            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.116        feature_extractor ([`CLIPImageProcessor`]):117            Model that extracts features from generated images to be used as inputs for the `safety_checker`.118    """119 120    _optional_components = ["safety_checker", "feature_extractor"]121 122    def __init__(123        self,124        vae: AutoencoderKL,125        text_encoder: CLIPTextModel,126        tokenizer: CLIPTokenizer,127        unet: UNet2DConditionModel,128        scheduler: KarrasDiffusionSchedulers,129        safety_checker: StableDiffusionSafetyChecker,130        feature_extractor: CLIPImageProcessor,131        requires_safety_checker: bool = True,132    ):133        super().__init__()134 135        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:136            deprecation_message = (137                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"138                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "139                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"140                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"141                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"142                " file"143            )144            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)145            new_config = dict(scheduler.config)146            new_config["steps_offset"] = 1147            scheduler._internal_dict = FrozenDict(new_config)148 149        if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:150            deprecation_message = (151                f"The configuration file of this scheduler: {scheduler} has not set the configuration"152                " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"153                " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"154                " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"155                " Hub, it would be very nice if you could open a Pull request for the"156                " `scheduler/scheduler_config.json` file"157            )158            deprecate(159                "skip_prk_steps not set",160                "1.0.0",161                deprecation_message,162                standard_warn=False,163            )164            new_config = dict(scheduler.config)165            new_config["skip_prk_steps"] = True166            scheduler._internal_dict = FrozenDict(new_config)167 168        if safety_checker is None and requires_safety_checker:169            logger.warning(170                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"171                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"172                " results in services or applications open to the public. Both the diffusers team and Hugging Face"173                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"174                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"175                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."176            )177 178        if safety_checker is not None and feature_extractor is None:179            raise ValueError(180                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"181                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."182            )183 184        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(185            version.parse(unet.config._diffusers_version).base_version186        ) < version.parse("0.9.0.dev0")187        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64188        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:189            deprecation_message = (190                "The configuration file of the unet has set the default `sample_size` to smaller than"191                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"192                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"193                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"194                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"195                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"196                " in the config might lead to incorrect results in future versions. If you have downloaded this"197                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"198                " the `unet/config.json` file"199            )200            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)201            new_config = dict(unet.config)202            new_config["sample_size"] = 64203            unet._internal_dict = FrozenDict(new_config)204        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4205        if unet.config.in_channels != 4:206            logger.warning(207                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"208                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"209                ". If you did not intend to modify"210                " this behavior, please check whether you have loaded the right checkpoint."211            )212 213        self.register_modules(214            vae=vae,215            text_encoder=text_encoder,216            tokenizer=tokenizer,217            unet=unet,218            scheduler=scheduler,219            safety_checker=safety_checker,220            feature_extractor=feature_extractor,221        )222        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)223        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)224        self.register_to_config(requires_safety_checker=requires_safety_checker)225 226    def _default_height_width(227        self,228        height: Optional[int],229        width: Optional[int],230        image: Union[PIL.Image.Image, torch.Tensor, List[PIL.Image.Image]],231    ) -> Tuple[int, int]:232        r"""233        Calculate the default height and width for the given image.234 235        Args:236            height (int or None): The desired height of the image. If None, the height will be determined based on the input image.237            width (int or None): The desired width of the image. If None, the width will be determined based on the input image.238            image (PIL.Image.Image or torch.Tensor or list[PIL.Image.Image]): The input image or a list of images.239 240        Returns:241            Tuple[int, int]: A tuple containing the calculated height and width.242 243        """244        # NOTE: It is possible that a list of images have different245        # dimensions for each image, so just checking the first image246        # is not _exactly_ correct, but it is simple.247        while isinstance(image, list):248            image = image[0]249 250        if height is None:251            if isinstance(image, PIL.Image.Image):252                height = image.height253            elif isinstance(image, torch.Tensor):254                height = image.shape[2]255 256            height = (height // 8) * 8  # round down to nearest multiple of 8257 258        if width is None:259            if isinstance(image, PIL.Image.Image):260                width = image.width261            elif isinstance(image, torch.Tensor):262                width = image.shape[3]263 264            width = (width // 8) * 8  # round down to nearest multiple of 8265 266        return height, width267 268    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs269    def check_inputs(270        self,271        prompt: Optional[Union[str, List[str]]],272        height: int,273        width: int,274        callback_steps: Optional[int],275        negative_prompt: Optional[str] = None,276        prompt_embeds: Optional[torch.Tensor] = None,277        negative_prompt_embeds: Optional[torch.Tensor] = None,278        ip_adapter_image: Optional[torch.Tensor] = None,279        ip_adapter_image_embeds: Optional[torch.Tensor] = None,280        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,281    ) -> None:282        """283        Check the validity of the input arguments for the diffusion model.284 285        Args:286            prompt (Optional[Union[str, List[str]]]): The prompt text or list of prompt texts.287            height (int): The height of the input image.288            width (int): The width of the input image.289            callback_steps (Optional[int]): The number of steps to perform the callback on.290            negative_prompt (Optional[str]): The negative prompt text.291            prompt_embeds (Optional[torch.Tensor]): The prompt embeddings.292            negative_prompt_embeds (Optional[torch.Tensor]): The negative prompt embeddings.293            ip_adapter_image (Optional[torch.Tensor]): The input adapter image.294            ip_adapter_image_embeds (Optional[torch.Tensor]): The input adapter image embeddings.295            callback_on_step_end_tensor_inputs (Optional[List[str]]): The list of tensor inputs to perform the callback on.296 297        Raises:298            ValueError: If `height` or `width` is not divisible by 8.299            ValueError: If `callback_steps` is not a positive integer.300            ValueError: If `callback_on_step_end_tensor_inputs` contains invalid tensor inputs.301            ValueError: If both `prompt` and `prompt_embeds` are provided.302            ValueError: If neither `prompt` nor `prompt_embeds` are provided.303            ValueError: If `prompt` is not of type `str` or `list`.304            ValueError: If both `negative_prompt` and `negative_prompt_embeds` are provided.305            ValueError: If both `prompt_embeds` and `negative_prompt_embeds` are provided and have different shapes.306            ValueError: If both `ip_adapter_image` and `ip_adapter_image_embeds` are provided.307 308        Returns:309            None310        """311        if height % 8 != 0 or width % 8 != 0:312            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")313 314        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):315            raise ValueError(316                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"317                f" {type(callback_steps)}."318            )319        if callback_on_step_end_tensor_inputs is not None and not all(320            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs321        ):322            raise ValueError(323                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"324            )325 326        if prompt is not None and prompt_embeds is not None:327            raise ValueError(328                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"329                " only forward one of the two."330            )331        elif prompt is None and prompt_embeds is None:332            raise ValueError(333                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."334            )335        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):336            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")337 338        if negative_prompt is not None and negative_prompt_embeds is not None:339            raise ValueError(340                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"341                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."342            )343 344        if prompt_embeds is not None and negative_prompt_embeds is not None:345            if prompt_embeds.shape != negative_prompt_embeds.shape:346                raise ValueError(347                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"348                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"349                    f" {negative_prompt_embeds.shape}."350                )351 352        if ip_adapter_image is not None and ip_adapter_image_embeds is not None:353            raise ValueError(354                "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."355            )356 357    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt358    def _encode_prompt(359        self,360        prompt: Union[str, List[str]],361        device: torch.device,362        num_images_per_prompt: int,363        do_classifier_free_guidance: bool,364        negative_prompt: Optional[Union[str, List[str]]] = None,365        prompt_embeds: Optional[torch.Tensor] = None,366        negative_prompt_embeds: Optional[torch.Tensor] = None,367        lora_scale: Optional[float] = None,368        **kwargs,369    ) -> torch.Tensor:370        r"""371        Encodes the prompt into embeddings.372 373        Args:374            prompt (Union[str, List[str]]): The prompt text or a list of prompt texts.375            device (torch.device): The device to use for encoding.376            num_images_per_prompt (int): The number of images per prompt.377            do_classifier_free_guidance (bool): Whether to use classifier-free guidance.378            negative_prompt (Optional[Union[str, List[str]]], optional): The negative prompt text or a list of negative prompt texts. Defaults to None.379            prompt_embeds (Optional[torch.Tensor], optional): The prompt embeddings. Defaults to None.380            negative_prompt_embeds (Optional[torch.Tensor], optional): The negative prompt embeddings. Defaults to None.381            lora_scale (Optional[float], optional): The LoRA scale. Defaults to None.382            **kwargs: Additional keyword arguments.383 384        Returns:385            torch.Tensor: The encoded prompt embeddings.386        """387        deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."388        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)389 390        prompt_embeds_tuple = self.encode_prompt(391            prompt=prompt,392            device=device,393            num_images_per_prompt=num_images_per_prompt,394            do_classifier_free_guidance=do_classifier_free_guidance,395            negative_prompt=negative_prompt,396            prompt_embeds=prompt_embeds,397            negative_prompt_embeds=negative_prompt_embeds,398            lora_scale=lora_scale,399            **kwargs,400        )401 402        # concatenate for backwards comp403        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])404 405        return prompt_embeds406 407    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt408    def encode_prompt(409        self,410        prompt: Optional[str],411        device: torch.device,412        num_images_per_prompt: int,413        do_classifier_free_guidance: bool,414        negative_prompt: Optional[str] = None,415        prompt_embeds: Optional[torch.Tensor] = None,416        negative_prompt_embeds: Optional[torch.Tensor] = None,417        lora_scale: Optional[float] = None,418        clip_skip: Optional[int] = None,419    ) -> torch.Tensor:420        r"""421        Encodes the prompt into text encoder hidden states.422 423        Args:424            prompt (`str` or `List[str]`, *optional*):425                prompt to be encoded426            device: (`torch.device`):427                torch device428            num_images_per_prompt (`int`):429                number of images that should be generated per prompt430            do_classifier_free_guidance (`bool`):431                whether to use classifier free guidance or not432            negative_prompt (`str` or `List[str]`, *optional*):433                The prompt or prompts not to guide the image generation. If not defined, one has to pass434                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is435                less than `1`).436            prompt_embeds (`torch.Tensor`, *optional*):437                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not438                provided, text embeddings will be generated from `prompt` input argument.439            negative_prompt_embeds (`torch.Tensor`, *optional*):440                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt441                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input442                argument.443            lora_scale (`float`, *optional*):444                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.445            clip_skip (`int`, *optional*):446                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that447                the output of the pre-final layer will be used for computing the prompt embeddings.448        """449        # set lora scale so that monkey patched LoRA450        # function of text encoder can correctly access it451        if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin):452            self._lora_scale = lora_scale453 454            # dynamically adjust the LoRA scale455            if not USE_PEFT_BACKEND:456                adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)457            else:458                scale_lora_layers(self.text_encoder, lora_scale)459 460        if prompt is not None and isinstance(prompt, str):461            batch_size = 1462        elif prompt is not None and isinstance(prompt, list):463            batch_size = len(prompt)464        else:465            batch_size = prompt_embeds.shape[0]466 467        if prompt_embeds is None:468            # textual inversion: process multi-vector tokens if necessary469            if isinstance(self, TextualInversionLoaderMixin):470                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)471 472            text_inputs = self.tokenizer(473                prompt,474                padding="max_length",475                max_length=self.tokenizer.model_max_length,476                truncation=True,477                return_tensors="pt",478            )479            text_input_ids = text_inputs.input_ids480            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids481 482            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(483                text_input_ids, untruncated_ids484            ):485                removed_text = self.tokenizer.batch_decode(486                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]487                )488                logger.warning(489                    "The following part of your input was truncated because CLIP can only handle sequences up to"490                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"491                )492 493            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:494                attention_mask = text_inputs.attention_mask.to(device)495            else:496                attention_mask = None497 498            if clip_skip is None:499                prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)500                prompt_embeds = prompt_embeds[0]501            else:502                prompt_embeds = self.text_encoder(503                    text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True504                )505                # Access the `hidden_states` first, that contains a tuple of506                # all the hidden states from the encoder layers. Then index into507                # the tuple to access the hidden states from the desired layer.508                prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]509                # We also need to apply the final LayerNorm here to not mess with the510                # representations. The `last_hidden_states` that we typically use for511                # obtaining the final prompt representations passes through the LayerNorm512                # layer.513                prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)514 515        if self.text_encoder is not None:516            prompt_embeds_dtype = self.text_encoder.dtype517        elif self.unet is not None:518            prompt_embeds_dtype = self.unet.dtype519        else:520            prompt_embeds_dtype = prompt_embeds.dtype521 522        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)523 524        bs_embed, seq_len, _ = prompt_embeds.shape525        # duplicate text embeddings for each generation per prompt, using mps friendly method526        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)527        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)528 529        # get unconditional embeddings for classifier free guidance530        if do_classifier_free_guidance and negative_prompt_embeds is None:531            uncond_tokens: List[str]532            if negative_prompt is None:533                uncond_tokens = [""] * batch_size534            elif prompt is not None and type(prompt) is not type(negative_prompt):535                raise TypeError(536                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="537                    f" {type(prompt)}."538                )539            elif isinstance(negative_prompt, str):540                uncond_tokens = [negative_prompt]541            elif batch_size != len(negative_prompt):542                raise ValueError(543                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"544                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"545                    " the batch size of `prompt`."546                )547            else:548                uncond_tokens = negative_prompt549 550            # textual inversion: process multi-vector tokens if necessary551            if isinstance(self, TextualInversionLoaderMixin):552                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)553 554            max_length = prompt_embeds.shape[1]555            uncond_input = self.tokenizer(556                uncond_tokens,557                padding="max_length",558                max_length=max_length,559                truncation=True,560                return_tensors="pt",561            )562 563            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:564                attention_mask = uncond_input.attention_mask.to(device)565            else:566                attention_mask = None567 568            negative_prompt_embeds = self.text_encoder(569                uncond_input.input_ids.to(device),570                attention_mask=attention_mask,571            )572            negative_prompt_embeds = negative_prompt_embeds[0]573 574        if do_classifier_free_guidance:575            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method576            seq_len = negative_prompt_embeds.shape[1]577 578            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)579 580            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)581            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)582 583        if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND:584            # Retrieve the original scale by scaling back the LoRA layers585            unscale_lora_layers(self.text_encoder, lora_scale)586 587        return prompt_embeds, negative_prompt_embeds588 589    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents590    def prepare_latents(591        self,592        batch_size: int,593        num_channels_latents: int,594        height: int,595        width: int,596        dtype: torch.dtype,597        device: torch.device,598        generator: Union[torch.Generator, List[torch.Generator]],599        latents: Optional[torch.Tensor] = None,600    ) -> torch.Tensor:601        r"""602        Prepare the latent vectors for diffusion.603 604        Args:605            batch_size (int): The number of samples in the batch.606            num_channels_latents (int): The number of channels in the latent vectors.607            height (int): The height of the latent vectors.608            width (int): The width of the latent vectors.609            dtype (torch.dtype): The data type of the latent vectors.610            device (torch.device): The device to place the latent vectors on.611            generator (Union[torch.Generator, List[torch.Generator]]): The generator(s) to use for random number generation.612            latents (Optional[torch.Tensor]): The pre-existing latent vectors. If None, new latent vectors will be generated.613 614        Returns:615            torch.Tensor: The prepared latent vectors.616        """617        shape = (618            batch_size,619            num_channels_latents,620            int(height) // self.vae_scale_factor,621            int(width) // self.vae_scale_factor,622        )623        if isinstance(generator, list) and len(generator) != batch_size:624            raise ValueError(625                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"626                f" size of {batch_size}. Make sure the batch size matches the length of the generators."627            )628 629        if latents is None:630            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)631        else:632            latents = latents.to(device)633 634        # scale the initial noise by the standard deviation required by the scheduler635        latents = latents * self.scheduler.init_noise_sigma636        return latents637 638    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs639    def prepare_extra_step_kwargs(640        self, generator: Union[torch.Generator, List[torch.Generator]], eta: float641    ) -> Dict[str, Any]:642        r"""643        Prepare extra keyword arguments for the scheduler step.644 645        Args:646            generator (Union[torch.Generator, List[torch.Generator]]): The generator used for sampling.647            eta (float): The value of eta (η) used with the DDIMScheduler. Should be between 0 and 1.648 649        Returns:650            Dict[str, Any]: A dictionary containing the extra keyword arguments for the scheduler step.651        """652        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature653        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.654        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502655        # and should be between [0, 1]656 657        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())658        extra_step_kwargs = {}659        if accepts_eta:660            extra_step_kwargs["eta"] = eta661 662        # check if the scheduler accepts generator663        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())664        if accepts_generator:665            extra_step_kwargs["generator"] = generator666        return extra_step_kwargs667 668    def prepare_image(669        self,670        image: Union[torch.Tensor, PIL.Image.Image, List[Union[torch.Tensor, PIL.Image.Image]]],671        width: int,672        height: int,673        batch_size: int,674        num_images_per_prompt: int,675        device: torch.device,676        dtype: torch.dtype,677        do_classifier_free_guidance: bool = False,678        guess_mode: bool = False,679    ) -> torch.Tensor:680        r"""681        Prepares the input image for processing.682 683        Args:684            image (torch.Tensor or PIL.Image.Image or list): The input image(s).685            width (int): The desired width of the image.686            height (int): The desired height of the image.687            batch_size (int): The batch size for processing.688            num_images_per_prompt (int): The number of images per prompt.689            device (torch.device): The device to use for processing.690            dtype (torch.dtype): The data type of the image.691            do_classifier_free_guidance (bool, optional): Whether to perform classifier-free guidance. Defaults to False.692            guess_mode (bool, optional): Whether to use guess mode. Defaults to False.693 694        Returns:695            torch.Tensor: The prepared image for processing.696        """697        if not isinstance(image, torch.Tensor):698            if isinstance(image, PIL.Image.Image):699                image = [image]700 701            if isinstance(image[0], PIL.Image.Image):702                images = []703 704                for image_ in image:705                    image_ = image_.convert("RGB")706                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])707                    image_ = np.array(image_)708                    image_ = image_[None, :]709                    images.append(image_)710 711                image = images712 713                image = np.concatenate(image, axis=0)714                image = np.array(image).astype(np.float32) / 255.0715                image = (image - 0.5) / 0.5716                image = image.transpose(0, 3, 1, 2)717                image = torch.from_numpy(image)718            elif isinstance(image[0], torch.Tensor):719                image = torch.cat(image, dim=0)720 721        image_batch_size = image.shape[0]722 723        if image_batch_size == 1:724            repeat_by = batch_size725        else:726            # image batch size is the same as prompt batch size727            repeat_by = num_images_per_prompt728 729        image = image.repeat_interleave(repeat_by, dim=0)730 731        image = image.to(device=device, dtype=dtype)732 733        if do_classifier_free_guidance and not guess_mode:734            image = torch.cat([image] * 2)735 736        return image737 738    def prepare_ref_latents(739        self,740        refimage: torch.Tensor,741        batch_size: int,742        dtype: torch.dtype,743        device: torch.device,744        generator: Union[int, List[int]],745        do_classifier_free_guidance: bool,746    ) -> torch.Tensor:747        r"""748        Prepares reference latents for generating images.749 750        Args:751            refimage (torch.Tensor): The reference image.752            batch_size (int): The desired batch size.753            dtype (torch.dtype): The data type of the tensors.754            device (torch.device): The device to perform computations on.755            generator (int or list): The generator index or a list of generator indices.756            do_classifier_free_guidance (bool): Whether to use classifier-free guidance.757 758        Returns:759            torch.Tensor: The prepared reference latents.760        """761        refimage = refimage.to(device=device, dtype=dtype)762 763        # encode the mask image into latents space so we can concatenate it to the latents764        if isinstance(generator, list):765            ref_image_latents = [766                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])767                for i in range(batch_size)768            ]769            ref_image_latents = torch.cat(ref_image_latents, dim=0)770        else:771            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)772        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents773 774        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method775        if ref_image_latents.shape[0] < batch_size:776            if not batch_size % ref_image_latents.shape[0] == 0:777                raise ValueError(778                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"779                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."780                    " Make sure the number of images that you pass is divisible by the total requested batch size."781                )782            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)783 784        # aligning device to prevent device errors when concating it with the latent model input785        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)786        return ref_image_latents787 788    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker789    def run_safety_checker(790        self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype791    ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:792        r"""793        Runs the safety checker on the given image.794 795        Args:796            image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.797            device (torch.device): The device to run the safety checker on.798            dtype (torch.dtype): The data type of the input image.799 800        Returns:801            (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and802            a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.803        """804        if self.safety_checker is None:805            has_nsfw_concept = None806        else:807            if torch.is_tensor(image):808                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")809            else:810                feature_extractor_input = self.image_processor.numpy_to_pil(image)811            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)812            image, has_nsfw_concept = self.safety_checker(813                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)814            )815        return image, has_nsfw_concept816 817    @torch.no_grad()818    def __call__(819        self,820        prompt: Union[str, List[str]] = None,821        ref_image: Union[torch.Tensor, PIL.Image.Image] = None,822        height: Optional[int] = None,823        width: Optional[int] = None,824        num_inference_steps: int = 50,825        guidance_scale: float = 7.5,826        negative_prompt: Optional[Union[str, List[str]]] = None,827        num_images_per_prompt: Optional[int] = 1,828        eta: float = 0.0,829        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,830        latents: Optional[torch.Tensor] = None,831        prompt_embeds: Optional[torch.Tensor] = None,832        negative_prompt_embeds: Optional[torch.Tensor] = None,833        output_type: Optional[str] = "pil",834        return_dict: bool = True,835        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,836        callback_steps: int = 1,837        cross_attention_kwargs: Optional[Dict[str, Any]] = None,838        guidance_rescale: float = 0.0,839        attention_auto_machine_weight: float = 1.0,840        gn_auto_machine_weight: float = 1.0,841        style_fidelity: float = 0.5,842        reference_attn: bool = True,843        reference_adain: bool = True,844    ):845        r"""846        Function invoked when calling the pipeline for generation.847 848        Args:849            prompt (`str` or `List[str]`, *optional*):850                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.851                instead.852            ref_image (`torch.Tensor`, `PIL.Image.Image`):853                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If854                the type is specified as `torch.Tensor`, it is passed to Reference Control as is. `PIL.Image.Image` can855                also be accepted as an image.856            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):857                The height in pixels of the generated image.858            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):859                The width in pixels of the generated image.860            num_inference_steps (`int`, *optional*, defaults to 50):861                The number of denoising steps. More denoising steps usually lead to a higher quality image at the862                expense of slower inference.863            guidance_scale (`float`, *optional*, defaults to 7.5):864                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).865                `guidance_scale` is defined as `w` of equation 2. of [Imagen866                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >867                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,868                usually at the expense of lower image quality.869            negative_prompt (`str` or `List[str]`, *optional*):870                The prompt or prompts not to guide the image generation. If not defined, one has to pass871                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is872                less than `1`).873            num_images_per_prompt (`int`, *optional*, defaults to 1):874                The number of images to generate per prompt.875            eta (`float`, *optional*, defaults to 0.0):876                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to877                [`schedulers.DDIMScheduler`], will be ignored for others.878            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):879                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)880                to make generation deterministic.881            latents (`torch.Tensor`, *optional*):882                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image883                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents884                tensor will ge generated by sampling using the supplied random `generator`.885            prompt_embeds (`torch.Tensor`, *optional*):886                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not887                provided, text embeddings will be generated from `prompt` input argument.888            negative_prompt_embeds (`torch.Tensor`, *optional*):889                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt890                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input891                argument.892            output_type (`str`, *optional*, defaults to `"pil"`):893                The output format of the generate image. Choose between894                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.895            return_dict (`bool`, *optional*, defaults to `True`):896                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a897                plain tuple.898            callback (`Callable`, *optional*):899                A function that will be called every `callback_steps` steps during inference. The function will be900                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.901            callback_steps (`int`, *optional*, defaults to 1):902                The frequency at which the `callback` function will be called. If not specified, the callback will be903                called at every step.904            cross_attention_kwargs (`dict`, *optional*):905                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under906                `self.processor` in907                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).908            guidance_rescale (`float`, *optional*, defaults to 0.0):909                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are910                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of911                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).912                Guidance rescale factor should fix overexposure when using zero terminal SNR.913            attention_auto_machine_weight (`float`):914                Weight of using reference query for self attention's context.915                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.916            gn_auto_machine_weight (`float`):917                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.918            style_fidelity (`float`):919                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,920                elif style_fidelity=0.0, prompt more important, else balanced.921            reference_attn (`bool`):922                Whether to use reference query for self attention's context.923            reference_adain (`bool`):924                Whether to use reference adain.925 926        Examples:927 928        Returns:929            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:930            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.931            When returning a tuple, the first element is a list with the generated images, and the second element is a932            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"933            (nsfw) content, according to the `safety_checker`.934        """935        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."936 937        # 0. Default height and width to unet938        height, width = self._default_height_width(height, width, ref_image)939 940        # 1. Check inputs. Raise error if not correct941        self.check_inputs(942            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds943        )944 945        # 2. Define call parameters946        if prompt is not None and isinstance(prompt, str):947            batch_size = 1948        elif prompt is not None and isinstance(prompt, list):949            batch_size = len(prompt)950        else:951            batch_size = prompt_embeds.shape[0]952 953        device = self._execution_device954        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)955        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`956        # corresponds to doing no classifier free guidance.957        do_classifier_free_guidance = guidance_scale > 1.0958 959        # 3. Encode input prompt960        text_encoder_lora_scale = (961            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None962        )963        prompt_embeds = self._encode_prompt(964            prompt,965            device,966            num_images_per_prompt,967            do_classifier_free_guidance,968            negative_prompt,969            prompt_embeds=prompt_embeds,970            negative_prompt_embeds=negative_prompt_embeds,971            lora_scale=text_encoder_lora_scale,972        )973 974        # 4. Preprocess reference image975        ref_image = self.prepare_image(976            image=ref_image,977            width=width,978            height=height,979            batch_size=batch_size * num_images_per_prompt,980            num_images_per_prompt=num_images_per_prompt,981            device=device,982            dtype=prompt_embeds.dtype,983        )984 985        # 5. Prepare timesteps986        self.scheduler.set_timesteps(num_inference_steps, device=device)987        timesteps = self.scheduler.timesteps988 989        # 6. Prepare latent variables990        num_channels_latents = self.unet.config.in_channels991        latents = self.prepare_latents(992            batch_size * num_images_per_prompt,993            num_channels_latents,994            height,995            width,996            prompt_embeds.dtype,997            device,998            generator,999            latents,1000        )1001 1002        # 7. Prepare reference latent variables1003        ref_image_latents = self.prepare_ref_latents(1004            ref_image,1005            batch_size * num_images_per_prompt,1006            prompt_embeds.dtype,1007            device,1008            generator,1009            do_classifier_free_guidance,1010        )1011 1012        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline1013        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)1014 1015        # 9. Modify self attention and group norm1016        MODE = "write"1017        uc_mask = (1018            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)1019            .type_as(ref_image_latents)1020            .bool()1021        )1022 1023        def hacked_basic_transformer_inner_forward(1024            self,1025            hidden_states: torch.Tensor,1026            attention_mask: Optional[torch.Tensor] = None,1027            encoder_hidden_states: Optional[torch.Tensor] = None,1028            encoder_attention_mask: Optional[torch.Tensor] = None,1029            timestep: Optional[torch.LongTensor] = None,1030            cross_attention_kwargs: Dict[str, Any] = None,1031            class_labels: Optional[torch.LongTensor] = None,1032        ):1033            if self.use_ada_layer_norm:1034                norm_hidden_states = self.norm1(hidden_states, timestep)1035            elif self.use_ada_layer_norm_zero:1036                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(1037                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype1038                )1039            else:1040                norm_hidden_states = self.norm1(hidden_states)1041 1042            # 1. Self-Attention1043            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}1044            if self.only_cross_attention:1045                attn_output = self.attn1(1046                    norm_hidden_states,1047                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1048                    attention_mask=attention_mask,1049                    **cross_attention_kwargs,1050                )1051            else:1052                if MODE == "write":1053                    self.bank.append(norm_hidden_states.detach().clone())1054                    attn_output = self.attn1(1055                        norm_hidden_states,1056                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1057                        attention_mask=attention_mask,1058                        **cross_attention_kwargs,1059                    )1060                if MODE == "read":1061                    if attention_auto_machine_weight > self.attn_weight:1062                        attn_output_uc = self.attn1(1063                            norm_hidden_states,1064                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),1065                            # attention_mask=attention_mask,1066                            **cross_attention_kwargs,1067                        )1068                        attn_output_c = attn_output_uc.clone()1069                        if do_classifier_free_guidance and style_fidelity > 0:1070                            attn_output_c[uc_mask] = self.attn1(1071                                norm_hidden_states[uc_mask],1072                                encoder_hidden_states=norm_hidden_states[uc_mask],1073                                **cross_attention_kwargs,1074                            )1075                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc1076                        self.bank.clear()1077                    else:1078                        attn_output = self.attn1(1079                            norm_hidden_states,1080                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1081                            attention_mask=attention_mask,1082                            **cross_attention_kwargs,1083                        )1084            if self.use_ada_layer_norm_zero:1085                attn_output = gate_msa.unsqueeze(1) * attn_output1086            hidden_states = attn_output + hidden_states1087 1088            if self.attn2 is not None:1089                norm_hidden_states = (1090                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)1091                )1092 1093                # 2. Cross-Attention1094                attn_output = self.attn2(1095                    norm_hidden_states,1096                    encoder_hidden_states=encoder_hidden_states,1097                    attention_mask=encoder_attention_mask,1098                    **cross_attention_kwargs,1099                )1100                hidden_states = attn_output + hidden_states1101 1102            # 3. Feed-forward1103            norm_hidden_states = self.norm3(hidden_states)1104 1105            if self.use_ada_layer_norm_zero:1106                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]1107 1108            ff_output = self.ff(norm_hidden_states)1109 1110            if self.use_ada_layer_norm_zero:1111                ff_output = gate_mlp.unsqueeze(1) * ff_output1112 1113            hidden_states = ff_output + hidden_states1114 1115            return hidden_states1116 1117        def hacked_mid_forward(self, *args, **kwargs):1118            eps = 1e-61119            x = self.original_forward(*args, **kwargs)1120            if MODE == "write":1121                if gn_auto_machine_weight >= self.gn_weight:1122                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1123                    self.mean_bank.append(mean)1124                    self.var_bank.append(var)1125            if MODE == "read":1126                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1127                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1128                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51129                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))1130                    var_acc = sum(self.var_bank) / float(len(self.var_bank))1131                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51132                    x_uc = (((x - mean) / std) * std_acc) + mean_acc1133                    x_c = x_uc.clone()1134                    if do_classifier_free_guidance and style_fidelity > 0:1135                        x_c[uc_mask] = x[uc_mask]1136                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc1137                self.mean_bank = []1138                self.var_bank = []1139            return x1140 1141        def hack_CrossAttnDownBlock2D_forward(1142            self,1143            hidden_states: torch.Tensor,1144            temb: Optional[torch.Tensor] = None,1145            encoder_hidden_states: Optional[torch.Tensor] = None,1146            attention_mask: Optional[torch.Tensor] = None,1147            cross_attention_kwargs: Optional[Dict[str, Any]] = None,1148            encoder_attention_mask: Optional[torch.Tensor] = None,1149        ):1150            eps = 1e-61151 1152            # TODO(Patrick, William) - attention mask is not used1153            output_states = ()1154 1155            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):1156                hidden_states = resnet(hidden_states, temb)1157                hidden_states = attn(1158                    hidden_states,1159                    encoder_hidden_states=encoder_hidden_states,1160                    cross_attention_kwargs=cross_attention_kwargs,1161                    attention_mask=attention_mask,1162                    encoder_attention_mask=encoder_attention_mask,1163                    return_dict=False,1164                )[0]1165                if MODE == "write":1166                    if gn_auto_machine_weight >= self.gn_weight:1167                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1168                        self.mean_bank.append([mean])1169                        self.var_bank.append([var])1170                if MODE == "read":1171                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1172                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1173                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51174                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))1175                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))1176                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51177                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc1178                        hidden_states_c = hidden_states_uc.clone()1179                        if do_classifier_free_guidance and style_fidelity > 0:1180                            hidden_states_c[uc_mask] = hidden_states[uc_mask]1181                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc1182 1183                output_states = output_states + (hidden_states,)1184 1185            if MODE == "read":1186                self.mean_bank = []1187                self.var_bank = []1188 1189            if self.downsamplers is not None:1190                for downsampler in self.downsamplers:1191                    hidden_states = downsampler(hidden_states)1192 1193                output_states = output_states + (hidden_states,)1194 1195            return hidden_states, output_states1196 1197        def hacked_DownBlock2D_forward(1198            self,1199            hidden_states: torch.Tensor,1200            temb: Optional[torch.Tensor] = None,

Showing the first 1,200 of 1466 lines. Download the file for the rest.