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 scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 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 scheduler is not None and getattr(scheduler.config, "skip_prk_steps", True) 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 = (185            unet is not None186            and hasattr(unet.config, "_diffusers_version")187            and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0")188        )189        is_unet_sample_size_less_64 = (190            unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64191        )192        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:193            deprecation_message = (194                "The configuration file of the unet has set the default `sample_size` to smaller than"195                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"196                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"197                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"198                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"199                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"200                " in the config might lead to incorrect results in future versions. If you have downloaded this"201                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"202                " the `unet/config.json` file"203            )204            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)205            new_config = dict(unet.config)206            new_config["sample_size"] = 64207            unet._internal_dict = FrozenDict(new_config)208        # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4209        if unet is not None and unet.config.in_channels != 4:210            logger.warning(211                f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"212                f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"213                ". If you did not intend to modify"214                " this behavior, please check whether you have loaded the right checkpoint."215            )216 217        self.register_modules(218            vae=vae,219            text_encoder=text_encoder,220            tokenizer=tokenizer,221            unet=unet,222            scheduler=scheduler,223            safety_checker=safety_checker,224            feature_extractor=feature_extractor,225        )226        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8227        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)228        self.register_to_config(requires_safety_checker=requires_safety_checker)229 230    def _default_height_width(231        self,232        height: Optional[int],233        width: Optional[int],234        image: Union[PIL.Image.Image, torch.Tensor, List[PIL.Image.Image]],235    ) -> Tuple[int, int]:236        r"""237        Calculate the default height and width for the given image.238 239        Args:240            height (int or None): The desired height of the image. If None, the height will be determined based on the input image.241            width (int or None): The desired width of the image. If None, the width will be determined based on the input image.242            image (PIL.Image.Image or torch.Tensor or list[PIL.Image.Image]): The input image or a list of images.243 244        Returns:245            Tuple[int, int]: A tuple containing the calculated height and width.246 247        """248        # NOTE: It is possible that a list of images have different249        # dimensions for each image, so just checking the first image250        # is not _exactly_ correct, but it is simple.251        while isinstance(image, list):252            image = image[0]253 254        if height is None:255            if isinstance(image, PIL.Image.Image):256                height = image.height257            elif isinstance(image, torch.Tensor):258                height = image.shape[2]259 260            height = (height // 8) * 8  # round down to nearest multiple of 8261 262        if width is None:263            if isinstance(image, PIL.Image.Image):264                width = image.width265            elif isinstance(image, torch.Tensor):266                width = image.shape[3]267 268            width = (width // 8) * 8  # round down to nearest multiple of 8269 270        return height, width271 272    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs273    def check_inputs(274        self,275        prompt: Optional[Union[str, List[str]]],276        height: int,277        width: int,278        callback_steps: Optional[int],279        negative_prompt: Optional[str] = None,280        prompt_embeds: Optional[torch.Tensor] = None,281        negative_prompt_embeds: Optional[torch.Tensor] = None,282        ip_adapter_image: Optional[torch.Tensor] = None,283        ip_adapter_image_embeds: Optional[torch.Tensor] = None,284        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,285    ) -> None:286        """287        Check the validity of the input arguments for the diffusion model.288 289        Args:290            prompt (Optional[Union[str, List[str]]]): The prompt text or list of prompt texts.291            height (int): The height of the input image.292            width (int): The width of the input image.293            callback_steps (Optional[int]): The number of steps to perform the callback on.294            negative_prompt (Optional[str]): The negative prompt text.295            prompt_embeds (Optional[torch.Tensor]): The prompt embeddings.296            negative_prompt_embeds (Optional[torch.Tensor]): The negative prompt embeddings.297            ip_adapter_image (Optional[torch.Tensor]): The input adapter image.298            ip_adapter_image_embeds (Optional[torch.Tensor]): The input adapter image embeddings.299            callback_on_step_end_tensor_inputs (Optional[List[str]]): The list of tensor inputs to perform the callback on.300 301        Raises:302            ValueError: If `height` or `width` is not divisible by 8.303            ValueError: If `callback_steps` is not a positive integer.304            ValueError: If `callback_on_step_end_tensor_inputs` contains invalid tensor inputs.305            ValueError: If both `prompt` and `prompt_embeds` are provided.306            ValueError: If neither `prompt` nor `prompt_embeds` are provided.307            ValueError: If `prompt` is not of type `str` or `list`.308            ValueError: If both `negative_prompt` and `negative_prompt_embeds` are provided.309            ValueError: If both `prompt_embeds` and `negative_prompt_embeds` are provided and have different shapes.310            ValueError: If both `ip_adapter_image` and `ip_adapter_image_embeds` are provided.311 312        Returns:313            None314        """315        if height % 8 != 0 or width % 8 != 0:316            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")317 318        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):319            raise ValueError(320                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"321                f" {type(callback_steps)}."322            )323        if callback_on_step_end_tensor_inputs is not None and not all(324            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs325        ):326            raise ValueError(327                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]}"328            )329 330        if prompt is not None and prompt_embeds is not None:331            raise ValueError(332                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"333                " only forward one of the two."334            )335        elif prompt is None and prompt_embeds is None:336            raise ValueError(337                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."338            )339        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):340            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")341 342        if negative_prompt is not None and negative_prompt_embeds is not None:343            raise ValueError(344                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"345                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."346            )347 348        if prompt_embeds is not None and negative_prompt_embeds is not None:349            if prompt_embeds.shape != negative_prompt_embeds.shape:350                raise ValueError(351                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"352                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"353                    f" {negative_prompt_embeds.shape}."354                )355 356        if ip_adapter_image is not None and ip_adapter_image_embeds is not None:357            raise ValueError(358                "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."359            )360 361    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt362    def _encode_prompt(363        self,364        prompt: Union[str, List[str]],365        device: torch.device,366        num_images_per_prompt: int,367        do_classifier_free_guidance: bool,368        negative_prompt: Optional[Union[str, List[str]]] = None,369        prompt_embeds: Optional[torch.Tensor] = None,370        negative_prompt_embeds: Optional[torch.Tensor] = None,371        lora_scale: Optional[float] = None,372        **kwargs,373    ) -> torch.Tensor:374        r"""375        Encodes the prompt into embeddings.376 377        Args:378            prompt (Union[str, List[str]]): The prompt text or a list of prompt texts.379            device (torch.device): The device to use for encoding.380            num_images_per_prompt (int): The number of images per prompt.381            do_classifier_free_guidance (bool): Whether to use classifier-free guidance.382            negative_prompt (Optional[Union[str, List[str]]], optional): The negative prompt text or a list of negative prompt texts. Defaults to None.383            prompt_embeds (Optional[torch.Tensor], optional): The prompt embeddings. Defaults to None.384            negative_prompt_embeds (Optional[torch.Tensor], optional): The negative prompt embeddings. Defaults to None.385            lora_scale (Optional[float], optional): The LoRA scale. Defaults to None.386            **kwargs: Additional keyword arguments.387 388        Returns:389            torch.Tensor: The encoded prompt embeddings.390        """391        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."392        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)393 394        prompt_embeds_tuple = self.encode_prompt(395            prompt=prompt,396            device=device,397            num_images_per_prompt=num_images_per_prompt,398            do_classifier_free_guidance=do_classifier_free_guidance,399            negative_prompt=negative_prompt,400            prompt_embeds=prompt_embeds,401            negative_prompt_embeds=negative_prompt_embeds,402            lora_scale=lora_scale,403            **kwargs,404        )405 406        # concatenate for backwards comp407        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])408 409        return prompt_embeds410 411    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt412    def encode_prompt(413        self,414        prompt: Optional[str],415        device: torch.device,416        num_images_per_prompt: int,417        do_classifier_free_guidance: bool,418        negative_prompt: Optional[str] = None,419        prompt_embeds: Optional[torch.Tensor] = None,420        negative_prompt_embeds: Optional[torch.Tensor] = None,421        lora_scale: Optional[float] = None,422        clip_skip: Optional[int] = None,423    ) -> torch.Tensor:424        r"""425        Encodes the prompt into text encoder hidden states.426 427        Args:428            prompt (`str` or `List[str]`, *optional*):429                prompt to be encoded430            device: (`torch.device`):431                torch device432            num_images_per_prompt (`int`):433                number of images that should be generated per prompt434            do_classifier_free_guidance (`bool`):435                whether to use classifier free guidance or not436            negative_prompt (`str` or `List[str]`, *optional*):437                The prompt or prompts not to guide the image generation. If not defined, one has to pass438                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is439                less than `1`).440            prompt_embeds (`torch.Tensor`, *optional*):441                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not442                provided, text embeddings will be generated from `prompt` input argument.443            negative_prompt_embeds (`torch.Tensor`, *optional*):444                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt445                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input446                argument.447            lora_scale (`float`, *optional*):448                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.449            clip_skip (`int`, *optional*):450                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that451                the output of the pre-final layer will be used for computing the prompt embeddings.452        """453        # set lora scale so that monkey patched LoRA454        # function of text encoder can correctly access it455        if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin):456            self._lora_scale = lora_scale457 458            # dynamically adjust the LoRA scale459            if not USE_PEFT_BACKEND:460                adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)461            else:462                scale_lora_layers(self.text_encoder, lora_scale)463 464        if prompt is not None and isinstance(prompt, str):465            batch_size = 1466        elif prompt is not None and isinstance(prompt, list):467            batch_size = len(prompt)468        else:469            batch_size = prompt_embeds.shape[0]470 471        if prompt_embeds is None:472            # textual inversion: process multi-vector tokens if necessary473            if isinstance(self, TextualInversionLoaderMixin):474                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)475 476            text_inputs = self.tokenizer(477                prompt,478                padding="max_length",479                max_length=self.tokenizer.model_max_length,480                truncation=True,481                return_tensors="pt",482            )483            text_input_ids = text_inputs.input_ids484            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids485 486            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(487                text_input_ids, untruncated_ids488            ):489                removed_text = self.tokenizer.batch_decode(490                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]491                )492                logger.warning(493                    "The following part of your input was truncated because CLIP can only handle sequences up to"494                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"495                )496 497            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:498                attention_mask = text_inputs.attention_mask.to(device)499            else:500                attention_mask = None501 502            if clip_skip is None:503                prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)504                prompt_embeds = prompt_embeds[0]505            else:506                prompt_embeds = self.text_encoder(507                    text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True508                )509                # Access the `hidden_states` first, that contains a tuple of510                # all the hidden states from the encoder layers. Then index into511                # the tuple to access the hidden states from the desired layer.512                prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]513                # We also need to apply the final LayerNorm here to not mess with the514                # representations. The `last_hidden_states` that we typically use for515                # obtaining the final prompt representations passes through the LayerNorm516                # layer.517                prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)518 519        if self.text_encoder is not None:520            prompt_embeds_dtype = self.text_encoder.dtype521        elif self.unet is not None:522            prompt_embeds_dtype = self.unet.dtype523        else:524            prompt_embeds_dtype = prompt_embeds.dtype525 526        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)527 528        bs_embed, seq_len, _ = prompt_embeds.shape529        # duplicate text embeddings for each generation per prompt, using mps friendly method530        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)531        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)532 533        # get unconditional embeddings for classifier free guidance534        if do_classifier_free_guidance and negative_prompt_embeds is None:535            uncond_tokens: List[str]536            if negative_prompt is None:537                uncond_tokens = [""] * batch_size538            elif prompt is not None and type(prompt) is not type(negative_prompt):539                raise TypeError(540                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="541                    f" {type(prompt)}."542                )543            elif isinstance(negative_prompt, str):544                uncond_tokens = [negative_prompt]545            elif batch_size != len(negative_prompt):546                raise ValueError(547                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"548                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"549                    " the batch size of `prompt`."550                )551            else:552                uncond_tokens = negative_prompt553 554            # textual inversion: process multi-vector tokens if necessary555            if isinstance(self, TextualInversionLoaderMixin):556                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)557 558            max_length = prompt_embeds.shape[1]559            uncond_input = self.tokenizer(560                uncond_tokens,561                padding="max_length",562                max_length=max_length,563                truncation=True,564                return_tensors="pt",565            )566 567            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:568                attention_mask = uncond_input.attention_mask.to(device)569            else:570                attention_mask = None571 572            negative_prompt_embeds = self.text_encoder(573                uncond_input.input_ids.to(device),574                attention_mask=attention_mask,575            )576            negative_prompt_embeds = negative_prompt_embeds[0]577 578        if do_classifier_free_guidance:579            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method580            seq_len = negative_prompt_embeds.shape[1]581 582            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)583 584            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)585            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)586 587        if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND:588            # Retrieve the original scale by scaling back the LoRA layers589            unscale_lora_layers(self.text_encoder, lora_scale)590 591        return prompt_embeds, negative_prompt_embeds592 593    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents594    def prepare_latents(595        self,596        batch_size: int,597        num_channels_latents: int,598        height: int,599        width: int,600        dtype: torch.dtype,601        device: torch.device,602        generator: Union[torch.Generator, List[torch.Generator]],603        latents: Optional[torch.Tensor] = None,604    ) -> torch.Tensor:605        r"""606        Prepare the latent vectors for diffusion.607 608        Args:609            batch_size (int): The number of samples in the batch.610            num_channels_latents (int): The number of channels in the latent vectors.611            height (int): The height of the latent vectors.612            width (int): The width of the latent vectors.613            dtype (torch.dtype): The data type of the latent vectors.614            device (torch.device): The device to place the latent vectors on.615            generator (Union[torch.Generator, List[torch.Generator]]): The generator(s) to use for random number generation.616            latents (Optional[torch.Tensor]): The pre-existing latent vectors. If None, new latent vectors will be generated.617 618        Returns:619            torch.Tensor: The prepared latent vectors.620        """621        shape = (622            batch_size,623            num_channels_latents,624            int(height) // self.vae_scale_factor,625            int(width) // self.vae_scale_factor,626        )627        if isinstance(generator, list) and len(generator) != batch_size:628            raise ValueError(629                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"630                f" size of {batch_size}. Make sure the batch size matches the length of the generators."631            )632 633        if latents is None:634            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)635        else:636            latents = latents.to(device)637 638        # scale the initial noise by the standard deviation required by the scheduler639        latents = latents * self.scheduler.init_noise_sigma640        return latents641 642    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs643    def prepare_extra_step_kwargs(644        self, generator: Union[torch.Generator, List[torch.Generator]], eta: float645    ) -> Dict[str, Any]:646        r"""647        Prepare extra keyword arguments for the scheduler step.648 649        Args:650            generator (Union[torch.Generator, List[torch.Generator]]): The generator used for sampling.651            eta (float): The value of eta (η) used with the DDIMScheduler. Should be between 0 and 1.652 653        Returns:654            Dict[str, Any]: A dictionary containing the extra keyword arguments for the scheduler step.655        """656        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature657        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.658        # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502659        # and should be between [0, 1]660 661        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())662        extra_step_kwargs = {}663        if accepts_eta:664            extra_step_kwargs["eta"] = eta665 666        # check if the scheduler accepts generator667        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())668        if accepts_generator:669            extra_step_kwargs["generator"] = generator670        return extra_step_kwargs671 672    def prepare_image(673        self,674        image: Union[torch.Tensor, PIL.Image.Image, List[Union[torch.Tensor, PIL.Image.Image]]],675        width: int,676        height: int,677        batch_size: int,678        num_images_per_prompt: int,679        device: torch.device,680        dtype: torch.dtype,681        do_classifier_free_guidance: bool = False,682        guess_mode: bool = False,683    ) -> torch.Tensor:684        r"""685        Prepares the input image for processing.686 687        Args:688            image (torch.Tensor or PIL.Image.Image or list): The input image(s).689            width (int): The desired width of the image.690            height (int): The desired height of the image.691            batch_size (int): The batch size for processing.692            num_images_per_prompt (int): The number of images per prompt.693            device (torch.device): The device to use for processing.694            dtype (torch.dtype): The data type of the image.695            do_classifier_free_guidance (bool, optional): Whether to perform classifier-free guidance. Defaults to False.696            guess_mode (bool, optional): Whether to use guess mode. Defaults to False.697 698        Returns:699            torch.Tensor: The prepared image for processing.700        """701        if not isinstance(image, torch.Tensor):702            if isinstance(image, PIL.Image.Image):703                image = [image]704 705            if isinstance(image[0], PIL.Image.Image):706                images = []707 708                for image_ in image:709                    image_ = image_.convert("RGB")710                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])711                    image_ = np.array(image_)712                    image_ = image_[None, :]713                    images.append(image_)714 715                image = images716 717                image = np.concatenate(image, axis=0)718                image = np.array(image).astype(np.float32) / 255.0719                image = (image - 0.5) / 0.5720                image = image.transpose(0, 3, 1, 2)721                image = torch.from_numpy(image)722            elif isinstance(image[0], torch.Tensor):723                image = torch.cat(image, dim=0)724 725        image_batch_size = image.shape[0]726 727        if image_batch_size == 1:728            repeat_by = batch_size729        else:730            # image batch size is the same as prompt batch size731            repeat_by = num_images_per_prompt732 733        image = image.repeat_interleave(repeat_by, dim=0)734 735        image = image.to(device=device, dtype=dtype)736 737        if do_classifier_free_guidance and not guess_mode:738            image = torch.cat([image] * 2)739 740        return image741 742    def prepare_ref_latents(743        self,744        refimage: torch.Tensor,745        batch_size: int,746        dtype: torch.dtype,747        device: torch.device,748        generator: Union[int, List[int]],749        do_classifier_free_guidance: bool,750    ) -> torch.Tensor:751        r"""752        Prepares reference latents for generating images.753 754        Args:755            refimage (torch.Tensor): The reference image.756            batch_size (int): The desired batch size.757            dtype (torch.dtype): The data type of the tensors.758            device (torch.device): The device to perform computations on.759            generator (int or list): The generator index or a list of generator indices.760            do_classifier_free_guidance (bool): Whether to use classifier-free guidance.761 762        Returns:763            torch.Tensor: The prepared reference latents.764        """765        refimage = refimage.to(device=device, dtype=dtype)766 767        # encode the mask image into latents space so we can concatenate it to the latents768        if isinstance(generator, list):769            ref_image_latents = [770                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])771                for i in range(batch_size)772            ]773            ref_image_latents = torch.cat(ref_image_latents, dim=0)774        else:775            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)776        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents777 778        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method779        if ref_image_latents.shape[0] < batch_size:780            if not batch_size % ref_image_latents.shape[0] == 0:781                raise ValueError(782                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"783                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."784                    " Make sure the number of images that you pass is divisible by the total requested batch size."785                )786            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)787 788        # aligning device to prevent device errors when concating it with the latent model input789        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)790        return ref_image_latents791 792    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker793    def run_safety_checker(794        self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype795    ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:796        r"""797        Runs the safety checker on the given image.798 799        Args:800            image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.801            device (torch.device): The device to run the safety checker on.802            dtype (torch.dtype): The data type of the input image.803 804        Returns:805            (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and806            a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.807        """808        if self.safety_checker is None:809            has_nsfw_concept = None810        else:811            if torch.is_tensor(image):812                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")813            else:814                feature_extractor_input = self.image_processor.numpy_to_pil(image)815            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)816            image, has_nsfw_concept = self.safety_checker(817                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)818            )819        return image, has_nsfw_concept820 821    @torch.no_grad()822    def __call__(823        self,824        prompt: Union[str, List[str]] = None,825        ref_image: Union[torch.Tensor, PIL.Image.Image] = None,826        height: Optional[int] = None,827        width: Optional[int] = None,828        num_inference_steps: int = 50,829        guidance_scale: float = 7.5,830        negative_prompt: Optional[Union[str, List[str]]] = None,831        num_images_per_prompt: Optional[int] = 1,832        eta: float = 0.0,833        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,834        latents: Optional[torch.Tensor] = None,835        prompt_embeds: Optional[torch.Tensor] = None,836        negative_prompt_embeds: Optional[torch.Tensor] = None,837        output_type: Optional[str] = "pil",838        return_dict: bool = True,839        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,840        callback_steps: int = 1,841        cross_attention_kwargs: Optional[Dict[str, Any]] = None,842        guidance_rescale: float = 0.0,843        attention_auto_machine_weight: float = 1.0,844        gn_auto_machine_weight: float = 1.0,845        style_fidelity: float = 0.5,846        reference_attn: bool = True,847        reference_adain: bool = True,848    ):849        r"""850        Function invoked when calling the pipeline for generation.851 852        Args:853            prompt (`str` or `List[str]`, *optional*):854                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.855                instead.856            ref_image (`torch.Tensor`, `PIL.Image.Image`):857                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If858                the type is specified as `torch.Tensor`, it is passed to Reference Control as is. `PIL.Image.Image` can859                also be accepted as an image.860            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):861                The height in pixels of the generated image.862            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):863                The width in pixels of the generated image.864            num_inference_steps (`int`, *optional*, defaults to 50):865                The number of denoising steps. More denoising steps usually lead to a higher quality image at the866                expense of slower inference.867            guidance_scale (`float`, *optional*, defaults to 7.5):868                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).869                `guidance_scale` is defined as `w` of equation 2. of [Imagen870                Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >871                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,872                usually at the expense of lower image quality.873            negative_prompt (`str` or `List[str]`, *optional*):874                The prompt or prompts not to guide the image generation. If not defined, one has to pass875                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is876                less than `1`).877            num_images_per_prompt (`int`, *optional*, defaults to 1):878                The number of images to generate per prompt.879            eta (`float`, *optional*, defaults to 0.0):880                Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to881                [`schedulers.DDIMScheduler`], will be ignored for others.882            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):883                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)884                to make generation deterministic.885            latents (`torch.Tensor`, *optional*):886                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image887                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents888                tensor will ge generated by sampling using the supplied random `generator`.889            prompt_embeds (`torch.Tensor`, *optional*):890                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not891                provided, text embeddings will be generated from `prompt` input argument.892            negative_prompt_embeds (`torch.Tensor`, *optional*):893                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt894                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input895                argument.896            output_type (`str`, *optional*, defaults to `"pil"`):897                The output format of the generate image. Choose between898                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.899            return_dict (`bool`, *optional*, defaults to `True`):900                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a901                plain tuple.902            callback (`Callable`, *optional*):903                A function that will be called every `callback_steps` steps during inference. The function will be904                called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.905            callback_steps (`int`, *optional*, defaults to 1):906                The frequency at which the `callback` function will be called. If not specified, the callback will be907                called at every step.908            cross_attention_kwargs (`dict`, *optional*):909                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under910                `self.processor` in911                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).912            guidance_rescale (`float`, *optional*, defaults to 0.0):913                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are914                Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of915                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891).916                Guidance rescale factor should fix overexposure when using zero terminal SNR.917            attention_auto_machine_weight (`float`):918                Weight of using reference query for self attention's context.919                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.920            gn_auto_machine_weight (`float`):921                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.922            style_fidelity (`float`):923                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,924                elif style_fidelity=0.0, prompt more important, else balanced.925            reference_attn (`bool`):926                Whether to use reference query for self attention's context.927            reference_adain (`bool`):928                Whether to use reference adain.929 930        Examples:931 932        Returns:933            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:934            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.935            When returning a tuple, the first element is a list with the generated images, and the second element is a936            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"937            (nsfw) content, according to the `safety_checker`.938        """939        assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."940 941        # 0. Default height and width to unet942        height, width = self._default_height_width(height, width, ref_image)943 944        # 1. Check inputs. Raise error if not correct945        self.check_inputs(946            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds947        )948 949        # 2. Define call parameters950        if prompt is not None and isinstance(prompt, str):951            batch_size = 1952        elif prompt is not None and isinstance(prompt, list):953            batch_size = len(prompt)954        else:955            batch_size = prompt_embeds.shape[0]956 957        device = self._execution_device958        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)959        # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`960        # corresponds to doing no classifier free guidance.961        do_classifier_free_guidance = guidance_scale > 1.0962 963        # 3. Encode input prompt964        text_encoder_lora_scale = (965            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None966        )967        prompt_embeds = self._encode_prompt(968            prompt,969            device,970            num_images_per_prompt,971            do_classifier_free_guidance,972            negative_prompt,973            prompt_embeds=prompt_embeds,974            negative_prompt_embeds=negative_prompt_embeds,975            lora_scale=text_encoder_lora_scale,976        )977 978        # 4. Preprocess reference image979        ref_image = self.prepare_image(980            image=ref_image,981            width=width,982            height=height,983            batch_size=batch_size * num_images_per_prompt,984            num_images_per_prompt=num_images_per_prompt,985            device=device,986            dtype=prompt_embeds.dtype,987        )988 989        # 5. Prepare timesteps990        self.scheduler.set_timesteps(num_inference_steps, device=device)991        timesteps = self.scheduler.timesteps992 993        # 6. Prepare latent variables994        num_channels_latents = self.unet.config.in_channels995        latents = self.prepare_latents(996            batch_size * num_images_per_prompt,997            num_channels_latents,998            height,999            width,1000            prompt_embeds.dtype,1001            device,1002            generator,1003            latents,1004        )1005 1006        # 7. Prepare reference latent variables1007        ref_image_latents = self.prepare_ref_latents(1008            ref_image,1009            batch_size * num_images_per_prompt,1010            prompt_embeds.dtype,1011            device,1012            generator,1013            do_classifier_free_guidance,1014        )1015 1016        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline1017        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)1018 1019        # 9. Modify self attention and group norm1020        MODE = "write"1021        uc_mask = (1022            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)1023            .type_as(ref_image_latents)1024            .bool()1025        )1026 1027        def hacked_basic_transformer_inner_forward(1028            self,1029            hidden_states: torch.Tensor,1030            attention_mask: Optional[torch.Tensor] = None,1031            encoder_hidden_states: Optional[torch.Tensor] = None,1032            encoder_attention_mask: Optional[torch.Tensor] = None,1033            timestep: Optional[torch.LongTensor] = None,1034            cross_attention_kwargs: Dict[str, Any] = None,1035            class_labels: Optional[torch.LongTensor] = None,1036        ):1037            if self.use_ada_layer_norm:1038                norm_hidden_states = self.norm1(hidden_states, timestep)1039            elif self.use_ada_layer_norm_zero:1040                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(1041                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype1042                )1043            else:1044                norm_hidden_states = self.norm1(hidden_states)1045 1046            # 1. Self-Attention1047            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}1048            if self.only_cross_attention:1049                attn_output = self.attn1(1050                    norm_hidden_states,1051                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1052                    attention_mask=attention_mask,1053                    **cross_attention_kwargs,1054                )1055            else:1056                if MODE == "write":1057                    self.bank.append(norm_hidden_states.detach().clone())1058                    attn_output = self.attn1(1059                        norm_hidden_states,1060                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1061                        attention_mask=attention_mask,1062                        **cross_attention_kwargs,1063                    )1064                if MODE == "read":1065                    if attention_auto_machine_weight > self.attn_weight:1066                        attn_output_uc = self.attn1(1067                            norm_hidden_states,1068                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),1069                            # attention_mask=attention_mask,1070                            **cross_attention_kwargs,1071                        )1072                        attn_output_c = attn_output_uc.clone()1073                        if do_classifier_free_guidance and style_fidelity > 0:1074                            attn_output_c[uc_mask] = self.attn1(1075                                norm_hidden_states[uc_mask],1076                                encoder_hidden_states=norm_hidden_states[uc_mask],1077                                **cross_attention_kwargs,1078                            )1079                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc1080                        self.bank.clear()1081                    else:1082                        attn_output = self.attn1(1083                            norm_hidden_states,1084                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1085                            attention_mask=attention_mask,1086                            **cross_attention_kwargs,1087                        )1088            if self.use_ada_layer_norm_zero:1089                attn_output = gate_msa.unsqueeze(1) * attn_output1090            hidden_states = attn_output + hidden_states1091 1092            if self.attn2 is not None:1093                norm_hidden_states = (1094                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)1095                )1096 1097                # 2. Cross-Attention1098                attn_output = self.attn2(1099                    norm_hidden_states,1100                    encoder_hidden_states=encoder_hidden_states,1101                    attention_mask=encoder_attention_mask,1102                    **cross_attention_kwargs,1103                )1104                hidden_states = attn_output + hidden_states1105 1106            # 3. Feed-forward1107            norm_hidden_states = self.norm3(hidden_states)1108 1109            if self.use_ada_layer_norm_zero:1110                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]1111 1112            ff_output = self.ff(norm_hidden_states)1113 1114            if self.use_ada_layer_norm_zero:1115                ff_output = gate_mlp.unsqueeze(1) * ff_output1116 1117            hidden_states = ff_output + hidden_states1118 1119            return hidden_states1120 1121        def hacked_mid_forward(self, *args, **kwargs):1122            eps = 1e-61123            x = self.original_forward(*args, **kwargs)1124            if MODE == "write":1125                if gn_auto_machine_weight >= self.gn_weight:1126                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1127                    self.mean_bank.append(mean)1128                    self.var_bank.append(var)1129            if MODE == "read":1130                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1131                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1132                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51133                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))1134                    var_acc = sum(self.var_bank) / float(len(self.var_bank))1135                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51136                    x_uc = (((x - mean) / std) * std_acc) + mean_acc1137                    x_c = x_uc.clone()1138                    if do_classifier_free_guidance and style_fidelity > 0:1139                        x_c[uc_mask] = x[uc_mask]1140                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc1141                self.mean_bank = []1142                self.var_bank = []1143            return x1144 1145        def hack_CrossAttnDownBlock2D_forward(1146            self,1147            hidden_states: torch.Tensor,1148            temb: Optional[torch.Tensor] = None,1149            encoder_hidden_states: Optional[torch.Tensor] = None,1150            attention_mask: Optional[torch.Tensor] = None,1151            cross_attention_kwargs: Optional[Dict[str, Any]] = None,1152            encoder_attention_mask: Optional[torch.Tensor] = None,1153        ):1154            eps = 1e-61155 1156            # TODO(Patrick, William) - attention mask is not used1157            output_states = ()1158 1159            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):1160                hidden_states = resnet(hidden_states, temb)1161                hidden_states = attn(1162                    hidden_states,1163                    encoder_hidden_states=encoder_hidden_states,1164                    cross_attention_kwargs=cross_attention_kwargs,1165                    attention_mask=attention_mask,1166                    encoder_attention_mask=encoder_attention_mask,1167                    return_dict=False,1168                )[0]1169                if MODE == "write":1170                    if gn_auto_machine_weight >= self.gn_weight:1171                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1172                        self.mean_bank.append([mean])1173                        self.var_bank.append([var])1174                if MODE == "read":1175                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1176                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1177                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51178                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))1179                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))1180                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51181                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc1182                        hidden_states_c = hidden_states_uc.clone()1183                        if do_classifier_free_guidance and style_fidelity > 0:1184                            hidden_states_c[uc_mask] = hidden_states[uc_mask]1185                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc1186 1187                output_states = output_states + (hidden_states,)1188 1189            if MODE == "read":1190                self.mean_bank = []1191                self.var_bank = []1192 1193            if self.downsamplers is not None:1194                for downsampler in self.downsamplers:1195                    hidden_states = downsampler(hidden_states)1196 1197                output_states = output_states + (hidden_states,)1198 1199            return hidden_states, output_states1200 

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