diffusers/community-pipelines-mirror

1# Copyright 2024 Jingyang Zhang and The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import abc16import inspect17import math18import numbers19from typing import Any, Callable, Dict, List, Optional, Union20 21import numpy as np22import torch23import torch.nn as nn24import torch.nn.functional as F25from packaging import version26from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection27 28from diffusers.configuration_utils import FrozenDict29from diffusers.image_processor import PipelineImageInput, VaeImageProcessor30from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin31from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel32from diffusers.models.attention_processor import Attention, FusedAttnProcessor2_033from diffusers.models.lora import adjust_lora_scale_text_encoder34from diffusers.pipelines.pipeline_utils import DiffusionPipeline35from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput36from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker37from diffusers.schedulers import KarrasDiffusionSchedulers38from diffusers.utils import (39    USE_PEFT_BACKEND,40    deprecate,41    logging,42    replace_example_docstring,43    scale_lora_layers,44    unscale_lora_layers,45)46from diffusers.utils.torch_utils import randn_tensor47 48 49logger = logging.get_logger(__name__)  # pylint: disable=invalid-name50 51EXAMPLE_DOC_STRING = """52    Examples:53        ```py54        >>> import torch55        >>> from diffusers import StableDiffusionPipeline56 57        >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)58        >>> pipe = pipe.to("cuda")59 60        >>> prompt = "a photo of an astronaut riding a horse on mars"61        >>> image = pipe(prompt).images[0]62        ```63"""64 65 66class GaussianSmoothing(nn.Module):67    """68    Copied from official repo: https://github.com/showlab/BoxDiff/blob/master/utils/gaussian_smoothing.py69    Apply gaussian smoothing on a70    1d, 2d or 3d tensor. Filtering is performed seperately for each channel71    in the input using a depthwise convolution.72    Arguments:73        channels (int, sequence): Number of channels of the input tensors. Output will74            have this number of channels as well.75        kernel_size (int, sequence): Size of the gaussian kernel.76        sigma (float, sequence): Standard deviation of the gaussian kernel.77        dim (int, optional): The number of dimensions of the data.78            Default value is 2 (spatial).79    """80 81    def __init__(self, channels, kernel_size, sigma, dim=2):82        super(GaussianSmoothing, self).__init__()83        if isinstance(kernel_size, numbers.Number):84            kernel_size = [kernel_size] * dim85        if isinstance(sigma, numbers.Number):86            sigma = [sigma] * dim87 88        # The gaussian kernel is the product of the89        # gaussian function of each dimension.90        kernel = 191        meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size])92        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):93            mean = (size - 1) / 294            kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2))95 96        # Make sure sum of values in gaussian kernel equals 1.97        kernel = kernel / torch.sum(kernel)98 99        # Reshape to depthwise convolutional weight100        kernel = kernel.view(1, 1, *kernel.size())101        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))102 103        self.register_buffer("weight", kernel)104        self.groups = channels105 106        if dim == 1:107            self.conv = F.conv1d108        elif dim == 2:109            self.conv = F.conv2d110        elif dim == 3:111            self.conv = F.conv3d112        else:113            raise RuntimeError("Only 1, 2 and 3 dimensions are supported. Received {}.".format(dim))114 115    def forward(self, input):116        """117        Apply gaussian filter to input.118        Arguments:119            input (torch.Tensor): Input to apply gaussian filter on.120        Returns:121            filtered (torch.Tensor): Filtered output.122        """123        return self.conv(input, weight=self.weight.to(input.dtype), groups=self.groups)124 125 126class AttendExciteCrossAttnProcessor:127    def __init__(self, attnstore, place_in_unet):128        super().__init__()129        self.attnstore = attnstore130        self.place_in_unet = place_in_unet131 132    def __call__(133        self,134        attn: Attention,135        hidden_states: torch.FloatTensor,136        encoder_hidden_states: Optional[torch.FloatTensor] = None,137        attention_mask: Optional[torch.FloatTensor] = None,138    ) -> torch.Tensor:139        batch_size, sequence_length, _ = hidden_states.shape140        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size=1)141        query = attn.to_q(hidden_states)142 143        is_cross = encoder_hidden_states is not None144        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states145        key = attn.to_k(encoder_hidden_states)146        value = attn.to_v(encoder_hidden_states)147 148        query = attn.head_to_batch_dim(query)149        key = attn.head_to_batch_dim(key)150        value = attn.head_to_batch_dim(value)151 152        attention_probs = attn.get_attention_scores(query, key, attention_mask)153        self.attnstore(attention_probs, is_cross, self.place_in_unet)154 155        hidden_states = torch.bmm(attention_probs, value)156        hidden_states = attn.batch_to_head_dim(hidden_states)157 158        # linear proj159        hidden_states = attn.to_out[0](hidden_states)160        # dropout161        hidden_states = attn.to_out[1](hidden_states)162 163        return hidden_states164 165 166class AttentionControl(abc.ABC):167    def step_callback(self, x_t):168        return x_t169 170    def between_steps(self):171        return172 173    # @property174    # def num_uncond_att_layers(self):175    #     return 0176 177    @abc.abstractmethod178    def forward(self, attn, is_cross: bool, place_in_unet: str):179        raise NotImplementedError180 181    def __call__(self, attn, is_cross: bool, place_in_unet: str):182        if self.cur_att_layer >= self.num_uncond_att_layers:183            self.forward(attn, is_cross, place_in_unet)184        self.cur_att_layer += 1185        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:186            self.cur_att_layer = 0187            self.cur_step += 1188            self.between_steps()189 190    def reset(self):191        self.cur_step = 0192        self.cur_att_layer = 0193 194    def __init__(self):195        self.cur_step = 0196        self.num_att_layers = -1197        self.cur_att_layer = 0198 199 200class AttentionStore(AttentionControl):201    @staticmethod202    def get_empty_store():203        return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []}204 205    def forward(self, attn, is_cross: bool, place_in_unet: str):206        key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"207        if attn.shape[1] <= 32**2:  # avoid memory overhead208            self.step_store[key].append(attn)209        return attn210 211    def between_steps(self):212        self.attention_store = self.step_store213        if self.save_global_store:214            with torch.no_grad():215                if len(self.global_store) == 0:216                    self.global_store = self.step_store217                else:218                    for key in self.global_store:219                        for i in range(len(self.global_store[key])):220                            self.global_store[key][i] += self.step_store[key][i].detach()221        self.step_store = self.get_empty_store()222        self.step_store = self.get_empty_store()223 224    def get_average_attention(self):225        average_attention = self.attention_store226        return average_attention227 228    def get_average_global_attention(self):229        average_attention = {230            key: [item / self.cur_step for item in self.global_store[key]] for key in self.attention_store231        }232        return average_attention233 234    def reset(self):235        super(AttentionStore, self).reset()236        self.step_store = self.get_empty_store()237        self.attention_store = {}238        self.global_store = {}239 240    def __init__(self, save_global_store=False):241        """242        Initialize an empty AttentionStore243        :param step_index: used to visualize only a specific step in the diffusion process244        """245        super(AttentionStore, self).__init__()246        self.save_global_store = save_global_store247        self.step_store = self.get_empty_store()248        self.attention_store = {}249        self.global_store = {}250        self.curr_step_index = 0251        self.num_uncond_att_layers = 0252 253 254def aggregate_attention(255    attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int256) -> torch.Tensor:257    """Aggregates the attention across the different layers and heads at the specified resolution."""258    out = []259    attention_maps = attention_store.get_average_attention()260 261    # for k, v in attention_maps.items():262    #     for vv in v:263    #         print(vv.shape)264    # exit()265 266    num_pixels = res**2267    for location in from_where:268        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:269            if item.shape[1] == num_pixels:270                cross_maps = item.reshape(1, -1, res, res, item.shape[-1])[select]271                out.append(cross_maps)272    out = torch.cat(out, dim=0)273    out = out.sum(0) / out.shape[0]274    return out275 276 277def register_attention_control(model, controller):278    attn_procs = {}279    cross_att_count = 0280    for name in model.unet.attn_processors.keys():281        # cross_attention_dim = None if name.endswith("attn1.processor") else model.unet.config.cross_attention_dim282        if name.startswith("mid_block"):283            # hidden_size = model.unet.config.block_out_channels[-1]284            place_in_unet = "mid"285        elif name.startswith("up_blocks"):286            # block_id = int(name[len("up_blocks.")])287            # hidden_size = list(reversed(model.unet.config.block_out_channels))[block_id]288            place_in_unet = "up"289        elif name.startswith("down_blocks"):290            # block_id = int(name[len("down_blocks.")])291            # hidden_size = model.unet.config.block_out_channels[block_id]292            place_in_unet = "down"293        else:294            continue295 296        cross_att_count += 1297        attn_procs[name] = AttendExciteCrossAttnProcessor(attnstore=controller, place_in_unet=place_in_unet)298    model.unet.set_attn_processor(attn_procs)299    controller.num_att_layers = cross_att_count300 301 302def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):303    """304    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and305    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4306    """307    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)308    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)309    # rescale the results from guidance (fixes overexposure)310    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)311    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images312    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg313    return noise_cfg314 315 316def retrieve_timesteps(317    scheduler,318    num_inference_steps: Optional[int] = None,319    device: Optional[Union[str, torch.device]] = None,320    timesteps: Optional[List[int]] = None,321    **kwargs,322):323    """324    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles325    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.326 327    Args:328        scheduler (`SchedulerMixin`):329            The scheduler to get timesteps from.330        num_inference_steps (`int`):331            The number of diffusion steps used when generating samples with a pre-trained model. If used,332            `timesteps` must be `None`.333        device (`str` or `torch.device`, *optional*):334            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.335        timesteps (`List[int]`, *optional*):336                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default337                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`338                must be `None`.339 340    Returns:341        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the342        second element is the number of inference steps.343    """344    if timesteps is not None:345        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())346        if not accepts_timesteps:347            raise ValueError(348                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"349                f" timestep schedules. Please check whether you are using the correct scheduler."350            )351        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)352        timesteps = scheduler.timesteps353        num_inference_steps = len(timesteps)354    else:355        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)356        timesteps = scheduler.timesteps357    return timesteps, num_inference_steps358 359 360class StableDiffusionBoxDiffPipeline(361    DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin362):363    r"""364    Pipeline for text-to-image generation using Stable Diffusion with BoxDiff.365 366    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods367    implemented for all pipelines (downloading, saving, running on a particular device, etc.).368 369    The pipeline also inherits the following loading methods:370        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings371        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights372        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights373        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files374        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters375 376    Args:377        vae ([`AutoencoderKL`]):378            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.379        text_encoder ([`~transformers.CLIPTextModel`]):380            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).381        tokenizer ([`~transformers.CLIPTokenizer`]):382            A `CLIPTokenizer` to tokenize text.383        unet ([`UNet2DConditionModel`]):384            A `UNet2DConditionModel` to denoise the encoded image latents.385        scheduler ([`SchedulerMixin`]):386            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of387            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].388        safety_checker ([`StableDiffusionSafetyChecker`]):389            Classification module that estimates whether generated images could be considered offensive or harmful.390            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details391            about a model's potential harms.392        feature_extractor ([`~transformers.CLIPImageProcessor`]):393            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.394    """395 396    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"397    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]398    _exclude_from_cpu_offload = ["safety_checker"]399    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]400 401    def __init__(402        self,403        vae: AutoencoderKL,404        text_encoder: CLIPTextModel,405        tokenizer: CLIPTokenizer,406        unet: UNet2DConditionModel,407        scheduler: KarrasDiffusionSchedulers,408        safety_checker: StableDiffusionSafetyChecker,409        feature_extractor: CLIPImageProcessor,410        image_encoder: CLIPVisionModelWithProjection = None,411        requires_safety_checker: bool = True,412    ):413        super().__init__()414 415        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:416            deprecation_message = (417                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"418                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "419                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"420                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"421                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"422                " file"423            )424            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)425            new_config = dict(scheduler.config)426            new_config["steps_offset"] = 1427            scheduler._internal_dict = FrozenDict(new_config)428 429        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:430            deprecation_message = (431                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."432                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"433                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"434                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"435                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"436            )437            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)438            new_config = dict(scheduler.config)439            new_config["clip_sample"] = False440            scheduler._internal_dict = FrozenDict(new_config)441 442        if safety_checker is None and requires_safety_checker:443            logger.warning(444                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"445                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"446                " results in services or applications open to the public. Both the diffusers team and Hugging Face"447                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"448                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"449                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."450            )451 452        if safety_checker is not None and feature_extractor is None:453            raise ValueError(454                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"455                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."456            )457 458        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(459            version.parse(unet.config._diffusers_version).base_version460        ) < version.parse("0.9.0.dev0")461        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64462        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:463            deprecation_message = (464                "The configuration file of the unet has set the default `sample_size` to smaller than"465                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"466                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"467                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"468                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"469                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"470                " in the config might lead to incorrect results in future versions. If you have downloaded this"471                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"472                " the `unet/config.json` file"473            )474            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)475            new_config = dict(unet.config)476            new_config["sample_size"] = 64477            unet._internal_dict = FrozenDict(new_config)478 479        self.register_modules(480            vae=vae,481            text_encoder=text_encoder,482            tokenizer=tokenizer,483            unet=unet,484            scheduler=scheduler,485            safety_checker=safety_checker,486            feature_extractor=feature_extractor,487            image_encoder=image_encoder,488        )489        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)490        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)491        self.register_to_config(requires_safety_checker=requires_safety_checker)492 493    def enable_vae_slicing(self):494        r"""495        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to496        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.497        """498        self.vae.enable_slicing()499 500    def disable_vae_slicing(self):501        r"""502        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to503        computing decoding in one step.504        """505        self.vae.disable_slicing()506 507    def enable_vae_tiling(self):508        r"""509        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to510        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow511        processing larger images.512        """513        self.vae.enable_tiling()514 515    def disable_vae_tiling(self):516        r"""517        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to518        computing decoding in one step.519        """520        self.vae.disable_tiling()521 522    def _encode_prompt(523        self,524        prompt,525        device,526        num_images_per_prompt,527        do_classifier_free_guidance,528        negative_prompt=None,529        prompt_embeds: Optional[torch.FloatTensor] = None,530        negative_prompt_embeds: Optional[torch.FloatTensor] = None,531        lora_scale: Optional[float] = None,532        **kwargs,533    ):534        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."535        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)536 537        prompt_embeds_tuple = self.encode_prompt(538            prompt=prompt,539            device=device,540            num_images_per_prompt=num_images_per_prompt,541            do_classifier_free_guidance=do_classifier_free_guidance,542            negative_prompt=negative_prompt,543            prompt_embeds=prompt_embeds,544            negative_prompt_embeds=negative_prompt_embeds,545            lora_scale=lora_scale,546            **kwargs,547        )548 549        # concatenate for backwards comp550        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])551 552        return prompt_embeds553 554    def encode_prompt(555        self,556        prompt,557        device,558        num_images_per_prompt,559        do_classifier_free_guidance,560        negative_prompt=None,561        prompt_embeds: Optional[torch.FloatTensor] = None,562        negative_prompt_embeds: Optional[torch.FloatTensor] = None,563        lora_scale: Optional[float] = None,564        clip_skip: Optional[int] = None,565    ):566        r"""567        Encodes the prompt into text encoder hidden states.568 569        Args:570            prompt (`str` or `List[str]`, *optional*):571                prompt to be encoded572            device: (`torch.device`):573                torch device574            num_images_per_prompt (`int`):575                number of images that should be generated per prompt576            do_classifier_free_guidance (`bool`):577                whether to use classifier free guidance or not578            negative_prompt (`str` or `List[str]`, *optional*):579                The prompt or prompts not to guide the image generation. If not defined, one has to pass580                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is581                less than `1`).582            prompt_embeds (`torch.FloatTensor`, *optional*):583                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not584                provided, text embeddings will be generated from `prompt` input argument.585            negative_prompt_embeds (`torch.FloatTensor`, *optional*):586                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt587                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input588                argument.589            lora_scale (`float`, *optional*):590                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.591            clip_skip (`int`, *optional*):592                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that593                the output of the pre-final layer will be used for computing the prompt embeddings.594        """595        # set lora scale so that monkey patched LoRA596        # function of text encoder can correctly access it597        if lora_scale is not None and isinstance(self, LoraLoaderMixin):598            self._lora_scale = lora_scale599 600            # dynamically adjust the LoRA scale601            if not USE_PEFT_BACKEND:602                adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)603            else:604                scale_lora_layers(self.text_encoder, lora_scale)605 606        if prompt is not None and isinstance(prompt, str):607            batch_size = 1608        elif prompt is not None and isinstance(prompt, list):609            batch_size = len(prompt)610        else:611            batch_size = prompt_embeds.shape[0]612 613        if prompt_embeds is None:614            # textual inversion: procecss multi-vector tokens if necessary615            if isinstance(self, TextualInversionLoaderMixin):616                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)617 618            text_inputs = self.tokenizer(619                prompt,620                padding="max_length",621                max_length=self.tokenizer.model_max_length,622                truncation=True,623                return_tensors="pt",624            )625            text_input_ids = text_inputs.input_ids626            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids627 628            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(629                text_input_ids, untruncated_ids630            ):631                removed_text = self.tokenizer.batch_decode(632                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]633                )634                logger.warning(635                    "The following part of your input was truncated because CLIP can only handle sequences up to"636                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"637                )638 639            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:640                attention_mask = text_inputs.attention_mask.to(device)641            else:642                attention_mask = None643 644            if clip_skip is None:645                prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)646                prompt_embeds = prompt_embeds[0]647            else:648                prompt_embeds = self.text_encoder(649                    text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True650                )651                # Access the `hidden_states` first, that contains a tuple of652                # all the hidden states from the encoder layers. Then index into653                # the tuple to access the hidden states from the desired layer.654                prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]655                # We also need to apply the final LayerNorm here to not mess with the656                # representations. The `last_hidden_states` that we typically use for657                # obtaining the final prompt representations passes through the LayerNorm658                # layer.659                prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)660 661        if self.text_encoder is not None:662            prompt_embeds_dtype = self.text_encoder.dtype663        elif self.unet is not None:664            prompt_embeds_dtype = self.unet.dtype665        else:666            prompt_embeds_dtype = prompt_embeds.dtype667 668        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)669 670        bs_embed, seq_len, _ = prompt_embeds.shape671        # duplicate text embeddings for each generation per prompt, using mps friendly method672        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)673        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)674 675        # get unconditional embeddings for classifier free guidance676        if do_classifier_free_guidance and negative_prompt_embeds is None:677            uncond_tokens: List[str]678            if negative_prompt is None:679                uncond_tokens = [""] * batch_size680            elif prompt is not None and type(prompt) is not type(negative_prompt):681                raise TypeError(682                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="683                    f" {type(prompt)}."684                )685            elif isinstance(negative_prompt, str):686                uncond_tokens = [negative_prompt]687            elif batch_size != len(negative_prompt):688                raise ValueError(689                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"690                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"691                    " the batch size of `prompt`."692                )693            else:694                uncond_tokens = negative_prompt695 696            # textual inversion: procecss multi-vector tokens if necessary697            if isinstance(self, TextualInversionLoaderMixin):698                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)699 700            max_length = prompt_embeds.shape[1]701            uncond_input = self.tokenizer(702                uncond_tokens,703                padding="max_length",704                max_length=max_length,705                truncation=True,706                return_tensors="pt",707            )708 709            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:710                attention_mask = uncond_input.attention_mask.to(device)711            else:712                attention_mask = None713 714            negative_prompt_embeds = self.text_encoder(715                uncond_input.input_ids.to(device),716                attention_mask=attention_mask,717            )718            negative_prompt_embeds = negative_prompt_embeds[0]719 720        if do_classifier_free_guidance:721            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method722            seq_len = negative_prompt_embeds.shape[1]723 724            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)725 726            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)727            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)728 729        if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:730            # Retrieve the original scale by scaling back the LoRA layers731            unscale_lora_layers(self.text_encoder, lora_scale)732 733        return text_inputs, prompt_embeds, negative_prompt_embeds734 735    def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):736        dtype = next(self.image_encoder.parameters()).dtype737 738        if not isinstance(image, torch.Tensor):739            image = self.feature_extractor(image, return_tensors="pt").pixel_values740 741        image = image.to(device=device, dtype=dtype)742        if output_hidden_states:743            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]744            image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)745            uncond_image_enc_hidden_states = self.image_encoder(746                torch.zeros_like(image), output_hidden_states=True747            ).hidden_states[-2]748            uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(749                num_images_per_prompt, dim=0750            )751            return image_enc_hidden_states, uncond_image_enc_hidden_states752        else:753            image_embeds = self.image_encoder(image).image_embeds754            image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)755            uncond_image_embeds = torch.zeros_like(image_embeds)756 757            return image_embeds, uncond_image_embeds758 759    def run_safety_checker(self, image, device, dtype):760        if self.safety_checker is None:761            has_nsfw_concept = None762        else:763            if torch.is_tensor(image):764                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")765            else:766                feature_extractor_input = self.image_processor.numpy_to_pil(image)767            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)768            image, has_nsfw_concept = self.safety_checker(769                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)770            )771        return image, has_nsfw_concept772 773    def decode_latents(self, latents):774        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"775        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)776 777        latents = 1 / self.vae.config.scaling_factor * latents778        image = self.vae.decode(latents, return_dict=False)[0]779        image = (image / 2 + 0.5).clamp(0, 1)780        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16781        image = image.cpu().permute(0, 2, 3, 1).float().numpy()782        return image783 784    def prepare_extra_step_kwargs(self, generator, eta):785        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature786        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.787        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502788        # and should be between [0, 1]789 790        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())791        extra_step_kwargs = {}792        if accepts_eta:793            extra_step_kwargs["eta"] = eta794 795        # check if the scheduler accepts generator796        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())797        if accepts_generator:798            extra_step_kwargs["generator"] = generator799        return extra_step_kwargs800 801    def check_inputs(802        self,803        prompt,804        height,805        width,806        boxdiff_phrases,807        boxdiff_boxes,808        callback_steps,809        negative_prompt=None,810        prompt_embeds=None,811        negative_prompt_embeds=None,812        callback_on_step_end_tensor_inputs=None,813    ):814        if height % 8 != 0 or width % 8 != 0:815            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")816 817        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):818            raise ValueError(819                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"820                f" {type(callback_steps)}."821            )822        if callback_on_step_end_tensor_inputs is not None and not all(823            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs824        ):825            raise ValueError(826                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]}"827            )828 829        if prompt is not None and prompt_embeds is not None:830            raise ValueError(831                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"832                " only forward one of the two."833            )834        elif prompt is None and prompt_embeds is None:835            raise ValueError(836                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."837            )838        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):839            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")840 841        if negative_prompt is not None and negative_prompt_embeds is not None:842            raise ValueError(843                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"844                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."845            )846 847        if prompt_embeds is not None and negative_prompt_embeds is not None:848            if prompt_embeds.shape != negative_prompt_embeds.shape:849                raise ValueError(850                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"851                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"852                    f" {negative_prompt_embeds.shape}."853                )854 855        if boxdiff_phrases is not None or boxdiff_boxes is not None:856            if not (boxdiff_phrases is not None and boxdiff_boxes is not None):857                raise ValueError("Either both `boxdiff_phrases` and `boxdiff_boxes` must be passed or none of them.")858 859            if not isinstance(boxdiff_phrases, list) or not isinstance(boxdiff_boxes, list):860                raise ValueError("`boxdiff_phrases` and `boxdiff_boxes` must be lists.")861 862            if len(boxdiff_phrases) != len(boxdiff_boxes):863                raise ValueError(864                    "`boxdiff_phrases` and `boxdiff_boxes` must have the same length,"865                    f" got: `boxdiff_phrases` {len(boxdiff_phrases)} != `boxdiff_boxes`"866                    f" {len(boxdiff_boxes)}."867                )868 869    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):870        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)871        if isinstance(generator, list) and len(generator) != batch_size:872            raise ValueError(873                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"874                f" size of {batch_size}. Make sure the batch size matches the length of the generators."875            )876 877        if latents is None:878            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)879        else:880            latents = latents.to(device)881 882        # scale the initial noise by the standard deviation required by the scheduler883        latents = latents * self.scheduler.init_noise_sigma884        return latents885 886    def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):887        r"""Enables the FreeU mechanism as in https://arxiv.org/abs/2309.11497.888 889        The suffixes after the scaling factors represent the stages where they are being applied.890 891        Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of the values892        that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.893 894        Args:895            s1 (`float`):896                Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to897                mitigate "oversmoothing effect" in the enhanced denoising process.898            s2 (`float`):899                Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to900                mitigate "oversmoothing effect" in the enhanced denoising process.901            b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.902            b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.903        """904        if not hasattr(self, "unet"):905            raise ValueError("The pipeline must have `unet` for using FreeU.")906        self.unet.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2)907 908    def disable_freeu(self):909        """Disables the FreeU mechanism if enabled."""910        self.unet.disable_freeu()911 912    # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.fuse_qkv_projections913    def fuse_qkv_projections(self, unet: bool = True, vae: bool = True):914        """915        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,916        key, value) are fused. For cross-attention modules, key and value projection matrices are fused.917 918        <Tip warning={true}>919 920        This API is 🧪 experimental.921 922        </Tip>923 924        Args:925            unet (`bool`, defaults to `True`): To apply fusion on the UNet.926            vae (`bool`, defaults to `True`): To apply fusion on the VAE.927        """928        self.fusing_unet = False929        self.fusing_vae = False930 931        if unet:932            self.fusing_unet = True933            self.unet.fuse_qkv_projections()934            self.unet.set_attn_processor(FusedAttnProcessor2_0())935 936        if vae:937            if not isinstance(self.vae, AutoencoderKL):938                raise ValueError("`fuse_qkv_projections()` is only supported for the VAE of type `AutoencoderKL`.")939 940            self.fusing_vae = True941            self.vae.fuse_qkv_projections()942            self.vae.set_attn_processor(FusedAttnProcessor2_0())943 944    # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.unfuse_qkv_projections945    def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True):946        """Disable QKV projection fusion if enabled.947 948        <Tip warning={true}>949 950        This API is 🧪 experimental.951 952        </Tip>953 954        Args:955            unet (`bool`, defaults to `True`): To apply fusion on the UNet.956            vae (`bool`, defaults to `True`): To apply fusion on the VAE.957 958        """959        if unet:960            if not self.fusing_unet:961                logger.warning("The UNet was not initially fused for QKV projections. Doing nothing.")962            else:963                self.unet.unfuse_qkv_projections()964                self.fusing_unet = False965 966        if vae:967            if not self.fusing_vae:968                logger.warning("The VAE was not initially fused for QKV projections. Doing nothing.")969            else:970                self.vae.unfuse_qkv_projections()971                self.fusing_vae = False972 973    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding974    def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):975        """976        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298977 978        Args:979            timesteps (`torch.Tensor`):980                generate embedding vectors at these timesteps981            embedding_dim (`int`, *optional*, defaults to 512):982                dimension of the embeddings to generate983            dtype:984                data type of the generated embeddings985 986        Returns:987            `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`988        """989        assert len(w.shape) == 1990        w = w * 1000.0991 992        half_dim = embedding_dim // 2993        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)994        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)995        emb = w.to(dtype)[:, None] * emb[None, :]996        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)997        if embedding_dim % 2 == 1:  # zero pad998            emb = torch.nn.functional.pad(emb, (0, 1))999        assert emb.shape == (w.shape[0], embedding_dim)1000        return emb1001 1002    @property1003    def guidance_scale(self):1004        return self._guidance_scale1005 1006    @property1007    def guidance_rescale(self):1008        return self._guidance_rescale1009 1010    @property1011    def clip_skip(self):1012        return self._clip_skip1013 1014    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)1015    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`1016    # corresponds to doing no classifier free guidance.1017    @property1018    def do_classifier_free_guidance(self):1019        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None1020 1021    @property1022    def cross_attention_kwargs(self):1023        return self._cross_attention_kwargs1024 1025    @property1026    def num_timesteps(self):1027        return self._num_timesteps1028 1029    @property1030    def interrupt(self):1031        return self._interrupt1032 1033    def _compute_max_attention_per_index(1034        self,1035        attention_maps: torch.Tensor,1036        indices_to_alter: List[int],1037        smooth_attentions: bool = False,1038        sigma: float = 0.5,1039        kernel_size: int = 3,1040        normalize_eot: bool = False,1041        bboxes: List[int] = None,1042        L: int = 1,1043        P: float = 0.2,1044    ) -> List[torch.Tensor]:1045        """Computes the maximum attention value for each of the tokens we wish to alter."""1046        last_idx = -11047        if normalize_eot:1048            prompt = self.prompt1049            if isinstance(self.prompt, list):1050                prompt = self.prompt[0]1051            last_idx = len(self.tokenizer(prompt)["input_ids"]) - 11052        attention_for_text = attention_maps[:, :, 1:last_idx]1053        attention_for_text *= 1001054        attention_for_text = torch.nn.functional.softmax(attention_for_text, dim=-1)1055 1056        # Shift indices since we removed the first token "1:last_idx"1057        indices_to_alter = [index - 1 for index in indices_to_alter]1058 1059        # Extract the maximum values1060        max_indices_list_fg = []1061        max_indices_list_bg = []1062        dist_x = []1063        dist_y = []1064 1065        cnt = 01066        for i in indices_to_alter:1067            image = attention_for_text[:, :, i]1068 1069            # TODO1070            # box = [max(round(b / (512 / image.shape[0])), 0) for b in bboxes[cnt]]1071            # x1, y1, x2, y2 = box1072            H, W = image.shape1073            x1 = min(max(round(bboxes[cnt][0] * W), 0), W)1074            y1 = min(max(round(bboxes[cnt][1] * H), 0), H)1075            x2 = min(max(round(bboxes[cnt][2] * W), 0), W)1076            y2 = min(max(round(bboxes[cnt][3] * H), 0), H)1077            box = [x1, y1, x2, y2]1078            cnt += 11079 1080            # coordinates to masks1081            obj_mask = torch.zeros_like(image)1082            ones_mask = torch.ones([y2 - y1, x2 - x1], dtype=obj_mask.dtype).to(obj_mask.device)1083            obj_mask[y1:y2, x1:x2] = ones_mask1084            bg_mask = 1 - obj_mask1085 1086            if smooth_attentions:1087                smoothing = GaussianSmoothing(channels=1, kernel_size=kernel_size, sigma=sigma, dim=2).to(image.device)1088                input = F.pad(image.unsqueeze(0).unsqueeze(0), (1, 1, 1, 1), mode="reflect")1089                image = smoothing(input).squeeze(0).squeeze(0)1090 1091            # Inner-Box constraint1092            k = (obj_mask.sum() * P).long()1093            max_indices_list_fg.append((image * obj_mask).reshape(-1).topk(k)[0].mean())1094 1095            # Outer-Box constraint1096            k = (bg_mask.sum() * P).long()1097            max_indices_list_bg.append((image * bg_mask).reshape(-1).topk(k)[0].mean())1098 1099            # Corner Constraint1100            gt_proj_x = torch.max(obj_mask, dim=0)[0]1101            gt_proj_y = torch.max(obj_mask, dim=1)[0]1102            corner_mask_x = torch.zeros_like(gt_proj_x)1103            corner_mask_y = torch.zeros_like(gt_proj_y)1104 1105            # create gt according to the number config.L1106            N = gt_proj_x.shape[0]1107            corner_mask_x[max(box[0] - L, 0) : min(box[0] + L + 1, N)] = 1.01108            corner_mask_x[max(box[2] - L, 0) : min(box[2] + L + 1, N)] = 1.01109            corner_mask_y[max(box[1] - L, 0) : min(box[1] + L + 1, N)] = 1.01110            corner_mask_y[max(box[3] - L, 0) : min(box[3] + L + 1, N)] = 1.01111            dist_x.append((F.l1_loss(image.max(dim=0)[0], gt_proj_x, reduction="none") * corner_mask_x).mean())1112            dist_y.append((F.l1_loss(image.max(dim=1)[0], gt_proj_y, reduction="none") * corner_mask_y).mean())1113 1114        return max_indices_list_fg, max_indices_list_bg, dist_x, dist_y1115 1116    def _aggregate_and_get_max_attention_per_token(1117        self,1118        attention_store: AttentionStore,1119        indices_to_alter: List[int],1120        attention_res: int = 16,1121        smooth_attentions: bool = False,1122        sigma: float = 0.5,1123        kernel_size: int = 3,1124        normalize_eot: bool = False,1125        bboxes: List[int] = None,1126        L: int = 1,1127        P: float = 0.2,1128    ):1129        """Aggregates the attention for each token and computes the max activation value for each token to alter."""1130        attention_maps = aggregate_attention(1131            attention_store=attention_store,1132            res=attention_res,1133            from_where=("up", "down", "mid"),1134            is_cross=True,1135            select=0,1136        )1137        max_attention_per_index_fg, max_attention_per_index_bg, dist_x, dist_y = self._compute_max_attention_per_index(1138            attention_maps=attention_maps,1139            indices_to_alter=indices_to_alter,1140            smooth_attentions=smooth_attentions,1141            sigma=sigma,1142            kernel_size=kernel_size,1143            normalize_eot=normalize_eot,1144            bboxes=bboxes,1145            L=L,1146            P=P,1147        )1148        return max_attention_per_index_fg, max_attention_per_index_bg, dist_x, dist_y1149 1150    @staticmethod1151    def _compute_loss(1152        max_attention_per_index_fg: List[torch.Tensor],1153        max_attention_per_index_bg: List[torch.Tensor],1154        dist_x: List[torch.Tensor],1155        dist_y: List[torch.Tensor],1156        return_losses: bool = False,1157    ) -> torch.Tensor:1158        """Computes the attend-and-excite loss using the maximum attention value for each token."""1159        losses_fg = [max(0, 1.0 - curr_max) for curr_max in max_attention_per_index_fg]1160        losses_bg = [max(0, curr_max) for curr_max in max_attention_per_index_bg]1161        loss = sum(losses_fg) + sum(losses_bg) + sum(dist_x) + sum(dist_y)1162        if return_losses:1163            return max(losses_fg), losses_fg1164        else:1165            return max(losses_fg), loss1166 1167    @staticmethod1168    def _update_latent(latents: torch.Tensor, loss: torch.Tensor, step_size: float) -> torch.Tensor:1169        """Update the latent according to the computed loss."""1170        grad_cond = torch.autograd.grad(loss.requires_grad_(True), [latents], retain_graph=True)[0]1171        latents = latents - step_size * grad_cond1172        return latents1173 1174    def _perform_iterative_refinement_step(1175        self,1176        latents: torch.Tensor,1177        indices_to_alter: List[int],1178        loss_fg: torch.Tensor,1179        threshold: float,1180        text_embeddings: torch.Tensor,1181        text_input,1182        attention_store: AttentionStore,1183        step_size: float,1184        t: int,1185        attention_res: int = 16,1186        smooth_attentions: bool = True,1187        sigma: float = 0.5,1188        kernel_size: int = 3,1189        max_refinement_steps: int = 20,1190        normalize_eot: bool = False,1191        bboxes: List[int] = None,1192        L: int = 1,1193        P: float = 0.2,1194    ):1195        """1196        Performs the iterative latent refinement introduced in the paper. Here, we continuously update the latent1197        code according to our loss objective until the given threshold is reached for all tokens.1198        """1199        iteration = 01200        target_loss = max(0, 1.0 - threshold)

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