chore: import upstream snapshot with attribution

invokeai/backend/stable_diffusion/diffusion/__init__.py

@@ -0,0 +1,7 @@
+"""
+Initialization file for invokeai.models.diffusion
+"""
+
+from invokeai.backend.stable_diffusion.diffusion.shared_invokeai_diffusion import (
+    InvokeAIDiffuserComponent,  # noqa: F401
+)

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -0,0 +1,366 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+import torch
+
+from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData
+
+if TYPE_CHECKING:
+    from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+    from invokeai.backend.stable_diffusion.denoise_context import UNetKwargs
+
+
+@dataclass
+class BasicConditioningInfo:
+    """SD 1/2 text conditioning information produced by Compel."""
+
+    embeds: torch.Tensor
+
+    def to(self, device, dtype=None):
+        self.embeds = self.embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class SDXLConditioningInfo(BasicConditioningInfo):
+    """SDXL text conditioning information produced by Compel."""
+
+    pooled_embeds: torch.Tensor
+    add_time_ids: torch.Tensor
+
+    def to(self, device, dtype=None):
+        self.pooled_embeds = self.pooled_embeds.to(device=device, dtype=dtype)
+        self.add_time_ids = self.add_time_ids.to(device=device, dtype=dtype)
+        return super().to(device=device, dtype=dtype)
+
+
+@dataclass
+class FLUXConditioningInfo:
+    clip_embeds: torch.Tensor
+    t5_embeds: torch.Tensor
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.clip_embeds = self.clip_embeds.to(device=device, dtype=dtype)
+        self.t5_embeds = self.t5_embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class SD3ConditioningInfo:
+    clip_l_pooled_embeds: torch.Tensor
+    clip_l_embeds: torch.Tensor
+    clip_g_pooled_embeds: torch.Tensor
+    clip_g_embeds: torch.Tensor
+    t5_embeds: torch.Tensor | None
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.clip_l_pooled_embeds = self.clip_l_pooled_embeds.to(device=device, dtype=dtype)
+        self.clip_l_embeds = self.clip_l_embeds.to(device=device, dtype=dtype)
+        self.clip_g_pooled_embeds = self.clip_g_pooled_embeds.to(device=device, dtype=dtype)
+        self.clip_g_embeds = self.clip_g_embeds.to(device=device, dtype=dtype)
+        if self.t5_embeds is not None:
+            self.t5_embeds = self.t5_embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class CogView4ConditioningInfo:
+    glm_embeds: torch.Tensor
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.glm_embeds = self.glm_embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class ZImageConditioningInfo:
+    """Z-Image text conditioning information from Qwen3 text encoder."""
+
+    prompt_embeds: torch.Tensor
+    """Text embeddings from Qwen3 encoder. Shape: (batch_size, seq_len, hidden_size)."""
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.prompt_embeds = self.prompt_embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class QwenImageConditioningInfo:
+    """Qwen Image Edit conditioning information from Qwen2.5-VL encoder."""
+
+    prompt_embeds: torch.Tensor
+    """Text/image embeddings from Qwen2.5-VL encoder. Shape: (batch_size, seq_len, hidden_size)."""
+
+    prompt_embeds_mask: torch.Tensor | None = None
+    """Attention mask for prompt_embeds. Shape: (batch_size, seq_len). 1 for valid, 0 for padding."""
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.prompt_embeds = self.prompt_embeds.to(device=device, dtype=dtype)
+        if self.prompt_embeds_mask is not None:
+            self.prompt_embeds_mask = self.prompt_embeds_mask.to(device=device)
+        return self
+
+
+@dataclass
+class AnimaConditioningInfo:
+    """Anima text conditioning information from Qwen3 0.6B encoder + T5-XXL tokenizer.
+
+    Anima uses a dual-conditioning scheme where Qwen3 hidden states are combined
+    with T5-XXL token IDs inside the LLM Adapter (part of the transformer).
+    """
+
+    qwen3_embeds: torch.Tensor
+    """Qwen3 0.6B hidden states. Shape: (seq_len, hidden_size) where hidden_size=1024."""
+
+    t5xxl_ids: torch.Tensor
+    """T5-XXL token IDs. Shape: (seq_len,)."""
+
+    t5xxl_weights: Optional[torch.Tensor] = None
+    """Per-token weights for prompt weighting. Shape: (seq_len,). None means uniform weight."""
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.qwen3_embeds = self.qwen3_embeds.to(device=device, dtype=dtype)
+        self.t5xxl_ids = self.t5xxl_ids.to(device=device)
+        if self.t5xxl_weights is not None:
+            self.t5xxl_weights = self.t5xxl_weights.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class ConditioningFieldData:
+    # If you change this class, adding more types, you _must_ update the instantiation of ObjectSerializerDisk in
+    # invokeai/app/api/dependencies.py, adding the types to the list of safe globals. If you do not, torch will be
+    # unable to deserialize the object and will raise an error.
+    conditionings: (
+        List[BasicConditioningInfo]
+        | List[SDXLConditioningInfo]
+        | List[FLUXConditioningInfo]
+        | List[SD3ConditioningInfo]
+        | List[CogView4ConditioningInfo]
+        | List[ZImageConditioningInfo]
+        | List[QwenImageConditioningInfo]
+        | List[AnimaConditioningInfo]
+    )
+
+
+@dataclass
+class IPAdapterConditioningInfo:
+    cond_image_prompt_embeds: torch.Tensor
+    """IP-Adapter image encoder conditioning embeddings.
+    Shape: (num_images, num_tokens, encoding_dim).
+    """
+    uncond_image_prompt_embeds: torch.Tensor
+    """IP-Adapter image encoding embeddings to use for unconditional generation.
+    Shape: (num_images, num_tokens, encoding_dim).
+    """
+
+
+@dataclass
+class IPAdapterData:
+    """Data class for IP-Adapter configuration.
+
+    Attributes:
+        ip_adapter_model: The IP-Adapter model to use.
+        ip_adapter_conditioning: The IP-Adapter conditioning data.
+        mask: The mask to apply to the IP-Adapter conditioning.
+        target_blocks: List of target attention block names to apply IP-Adapter to.
+        negative_blocks: List of target attention block names that should use negative attention.
+        weight: The weight to apply to the IP-Adapter conditioning.
+        begin_step_percent: The percentage of steps at which to start applying the IP-Adapter.
+        end_step_percent: The percentage of steps at which to stop applying the IP-Adapter.
+        method: The method to use for applying the IP-Adapter ('full', 'style', 'composition').
+    """
+
+    ip_adapter_model: IPAdapter
+    ip_adapter_conditioning: IPAdapterConditioningInfo
+    mask: torch.Tensor
+    target_blocks: List[str]
+    negative_blocks: List[str] = field(default_factory=list)
+    weight: Union[float, List[float]] = 1.0
+    begin_step_percent: float = 0.0
+    end_step_percent: float = 1.0
+    method: str = "full"
+
+    def scale_for_step(self, step_index: int, total_steps: int) -> float:
+        first_adapter_step = math.floor(self.begin_step_percent * total_steps)
+        last_adapter_step = math.ceil(self.end_step_percent * total_steps)
+        weight = self.weight[step_index] if isinstance(self.weight, List) else self.weight
+        if step_index >= first_adapter_step and step_index <= last_adapter_step:
+            # Only apply this IP-Adapter if the current step is within the IP-Adapter's begin/end step range.
+            return weight
+        # Otherwise, set the IP-Adapter's scale to 0, so it has no effect.
+        return 0.0
+
+
+@dataclass
+class Range:
+    start: int
+    end: int
+
+
+class TextConditioningRegions:
+    def __init__(
+        self,
+        masks: torch.Tensor,
+        ranges: list[Range],
+    ):
+        # A binary mask indicating the regions of the image that the prompt should be applied to.
+        # Shape: (1, num_prompts, height, width)
+        # Dtype: torch.bool
+        self.masks = masks
+
+        # A list of ranges indicating the start and end indices of the embeddings that corresponding mask applies to.
+        # ranges[i] contains the embedding range for the i'th prompt / mask.
+        self.ranges = ranges
+
+        assert self.masks.shape[1] == len(self.ranges)
+
+
+class ConditioningMode(Enum):
+    Both = "both"
+    Negative = "negative"
+    Positive = "positive"
+
+
+class TextConditioningData:
+    def __init__(
+        self,
+        uncond_text: Union[BasicConditioningInfo, SDXLConditioningInfo],
+        cond_text: Union[BasicConditioningInfo, SDXLConditioningInfo],
+        uncond_regions: Optional[TextConditioningRegions],
+        cond_regions: Optional[TextConditioningRegions],
+        guidance_scale: Union[float, List[float]],
+        guidance_rescale_multiplier: float = 0,  # TODO: old backend, remove
+    ):
+        self.uncond_text = uncond_text
+        self.cond_text = cond_text
+        self.uncond_regions = uncond_regions
+        self.cond_regions = cond_regions
+        # Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+        # `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf).
+        # Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate
+        # images that are closely linked to the text `prompt`, usually at the expense of lower image quality.
+        self.guidance_scale = guidance_scale
+        # TODO: old backend, remove
+        # For models trained using zero-terminal SNR ("ztsnr"), it's suggested to use guidance_rescale_multiplier of 0.7.
+        # See [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+        self.guidance_rescale_multiplier = guidance_rescale_multiplier
+
+    def is_sdxl(self):
+        assert isinstance(self.uncond_text, SDXLConditioningInfo) == isinstance(self.cond_text, SDXLConditioningInfo)
+        return isinstance(self.cond_text, SDXLConditioningInfo)
+
+    def to_unet_kwargs(self, unet_kwargs: UNetKwargs, conditioning_mode: ConditioningMode):
+        """Fills unet arguments with data from provided conditionings.
+
+        Args:
+            unet_kwargs (UNetKwargs): Object which stores UNet model arguments.
+            conditioning_mode (ConditioningMode): Describes which conditionings should be used.
+        """
+        _, _, h, w = unet_kwargs.sample.shape
+        device = unet_kwargs.sample.device
+        dtype = unet_kwargs.sample.dtype
+
+        # TODO: combine regions with conditionings
+        if conditioning_mode == ConditioningMode.Both:
+            conditionings = [self.uncond_text, self.cond_text]
+            c_regions = [self.uncond_regions, self.cond_regions]
+        elif conditioning_mode == ConditioningMode.Positive:
+            conditionings = [self.cond_text]
+            c_regions = [self.cond_regions]
+        elif conditioning_mode == ConditioningMode.Negative:
+            conditionings = [self.uncond_text]
+            c_regions = [self.uncond_regions]
+        else:
+            raise ValueError(f"Unexpected conditioning mode: {conditioning_mode}")
+
+        encoder_hidden_states, encoder_attention_mask = self._concat_conditionings_for_batch(
+            [c.embeds for c in conditionings]
+        )
+
+        unet_kwargs.encoder_hidden_states = encoder_hidden_states
+        unet_kwargs.encoder_attention_mask = encoder_attention_mask
+
+        if self.is_sdxl():
+            added_cond_kwargs = dict(  # noqa: C408
+                text_embeds=torch.cat([c.pooled_embeds for c in conditionings]),
+                time_ids=torch.cat([c.add_time_ids for c in conditionings]),
+            )
+
+            unet_kwargs.added_cond_kwargs = added_cond_kwargs
+
+        if any(r is not None for r in c_regions):
+            tmp_regions = []
+            for c, r in zip(conditionings, c_regions, strict=True):
+                if r is None:
+                    r = TextConditioningRegions(
+                        masks=torch.ones((1, 1, h, w), dtype=dtype),
+                        ranges=[Range(start=0, end=c.embeds.shape[1])],
+                    )
+                tmp_regions.append(r)
+
+            if unet_kwargs.cross_attention_kwargs is None:
+                unet_kwargs.cross_attention_kwargs = {}
+
+            unet_kwargs.cross_attention_kwargs.update(
+                regional_prompt_data=RegionalPromptData(regions=tmp_regions, device=device, dtype=dtype),
+            )
+
+    @staticmethod
+    def _pad_zeros(t: torch.Tensor, pad_shape: tuple, dim: int) -> torch.Tensor:
+        return torch.cat([t, torch.zeros(pad_shape, device=t.device, dtype=t.dtype)], dim=dim)
+
+    @classmethod
+    def _pad_conditioning(
+        cls,
+        cond: torch.Tensor,
+        target_len: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Pad provided conditioning tensor to target_len by zeros and returns mask of unpadded bytes.
+
+        Args:
+            cond (torch.Tensor): Conditioning tensor which to pads by zeros.
+            target_len (int): To which length(tokens count) pad tensor.
+        """
+        conditioning_attention_mask = torch.ones((cond.shape[0], cond.shape[1]), device=cond.device, dtype=cond.dtype)
+
+        if cond.shape[1] < target_len:
+            conditioning_attention_mask = cls._pad_zeros(
+                conditioning_attention_mask,
+                pad_shape=(cond.shape[0], target_len - cond.shape[1]),
+                dim=1,
+            )
+
+            cond = cls._pad_zeros(
+                cond,
+                pad_shape=(cond.shape[0], target_len - cond.shape[1], cond.shape[2]),
+                dim=1,
+            )
+
+        return cond, conditioning_attention_mask
+
+    @classmethod
+    def _concat_conditionings_for_batch(
+        cls,
+        conditionings: List[torch.Tensor],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Concatenate provided conditioning tensors to one batched tensor.
+        If tensors have different sizes then pad them by zeros and creates
+        encoder_attention_mask to exclude padding from attention.
+
+        Args:
+            conditionings (List[torch.Tensor]): List of conditioning tensors to concatenate.
+        """
+        encoder_attention_mask = None
+        max_len = max([c.shape[1] for c in conditionings])
+        if any(c.shape[1] != max_len for c in conditionings):
+            encoder_attention_masks = [None] * len(conditionings)
+            for i in range(len(conditionings)):
+                conditionings[i], encoder_attention_masks[i] = cls._pad_conditioning(conditionings[i], max_len)
+            encoder_attention_mask = torch.cat(encoder_attention_masks)
+
+        return torch.cat(conditionings), encoder_attention_mask

invokeai/backend/stable_diffusion/diffusion/custom_atttention.py

@@ -0,0 +1,219 @@
+from dataclasses import dataclass
+from typing import List, Optional, cast
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention, AttnProcessor2_0
+
+from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionProcessorWeights
+from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData
+from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData
+
+
+@dataclass
+class IPAdapterAttentionWeights:
+    ip_adapter_weights: IPAttentionProcessorWeights
+    skip: bool
+    negative: bool
+
+
+class CustomAttnProcessor2_0(AttnProcessor2_0):
+    """A custom implementation of AttnProcessor2_0 that supports additional Invoke features.
+    This implementation is based on
+    https://github.com/huggingface/diffusers/blame/fcfa270fbd1dc294e2f3a505bae6bcb791d721c3/src/diffusers/models/attention_processor.py#L1204
+    Supported custom features:
+    - IP-Adapter
+    - Regional prompt attention
+    """
+
+    def __init__(
+        self,
+        ip_adapter_attention_weights: Optional[List[IPAdapterAttentionWeights]] = None,
+    ):
+        """Initialize a CustomAttnProcessor2_0.
+        Note: Arguments that are the same for all attention layers are passed to __call__(). Arguments that are
+        layer-specific are passed to __init__().
+        Args:
+            ip_adapter_weights: The IP-Adapter attention weights. ip_adapter_weights[i] contains the attention weights
+                for the i'th IP-Adapter.
+        """
+        super().__init__()
+        self._ip_adapter_attention_weights = ip_adapter_attention_weights
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        # For Regional Prompting:
+        regional_prompt_data: Optional[RegionalPromptData] = None,
+        percent_through: Optional[torch.Tensor] = None,
+        # For IP-Adapter:
+        regional_ip_data: Optional[RegionalIPData] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        """Apply attention.
+        Args:
+            regional_prompt_data: The regional prompt data for the current batch. If not None, this will be used to
+                apply regional prompt masking.
+            regional_ip_data: The IP-Adapter data for the current batch.
+        """
+        # If true, we are doing cross-attention, if false we are doing self-attention.
+        is_cross_attention = encoder_hidden_states is not None
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End unmodified block from AttnProcessor2_0.
+
+        _, query_seq_len, _ = hidden_states.shape
+        # Handle regional prompt attention masks.
+        if regional_prompt_data is not None and is_cross_attention:
+            assert percent_through is not None
+            prompt_region_attention_mask = regional_prompt_data.get_cross_attn_mask(
+                query_seq_len=query_seq_len, key_seq_len=sequence_length
+            )
+
+            if attention_mask is None:
+                attention_mask = prompt_region_attention_mask
+            else:
+                attention_mask = prompt_region_attention_mask + attention_mask
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End unmodified block from AttnProcessor2_0.
+
+        # Apply IP-Adapter conditioning.
+        if is_cross_attention:
+            if self._ip_adapter_attention_weights:
+                assert regional_ip_data is not None
+                ip_masks = regional_ip_data.get_masks(query_seq_len=query_seq_len)
+
+                assert (
+                    len(regional_ip_data.image_prompt_embeds)
+                    == len(self._ip_adapter_attention_weights)
+                    == len(regional_ip_data.scales)
+                    == ip_masks.shape[1]
+                )
+
+                for ipa_index, ipa_embed in enumerate(regional_ip_data.image_prompt_embeds):
+                    ipa_weights = self._ip_adapter_attention_weights[ipa_index].ip_adapter_weights
+                    ipa_scale = regional_ip_data.scales[ipa_index]
+                    ip_mask = ip_masks[0, ipa_index, ...]
+
+                    # The batch dimensions should match.
+                    assert ipa_embed.shape[0] == encoder_hidden_states.shape[0]
+                    # The token_len dimensions should match.
+                    assert ipa_embed.shape[-1] == encoder_hidden_states.shape[-1]
+
+                    ip_hidden_states = ipa_embed
+
+                    # Expected ip_hidden_state shape: (batch_size, num_ip_images, ip_seq_len, ip_image_embedding)
+
+                    if not self._ip_adapter_attention_weights[ipa_index].skip:
+                        # apply the IP-Adapter weights to the negative embeds
+                        if self._ip_adapter_attention_weights[ipa_index].negative:
+                            ip_hidden_states = torch.cat([ip_hidden_states[1], ip_hidden_states[0] * 0], dim=0)
+
+                        ip_key = ipa_weights.to_k_ip(ip_hidden_states)
+                        ip_value = ipa_weights.to_v_ip(ip_hidden_states)
+
+                        # Expected ip_key and ip_value shape:
+                        # (batch_size, num_ip_images, ip_seq_len, head_dim * num_heads)
+
+                        ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                        ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+                        # Expected ip_key and ip_value shape:
+                        # (batch_size, num_heads, num_ip_images * ip_seq_len, head_dim)
+
+                        # TODO: add support for attn.scale when we move to Torch 2.1
+                        ip_hidden_states = F.scaled_dot_product_attention(
+                            query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+                        )
+
+                        # Expected ip_hidden_states shape: (batch_size, num_heads, query_seq_len, head_dim)
+                        ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(
+                            batch_size, -1, attn.heads * head_dim
+                        )
+
+                        ip_hidden_states = ip_hidden_states.to(query.dtype)
+
+                        # Expected ip_hidden_states shape: (batch_size, query_seq_len, num_heads * head_dim)
+                        hidden_states = hidden_states + ipa_scale * ip_hidden_states * ip_mask
+            else:
+                # If IP-Adapter is not enabled, then regional_ip_data should not be passed in.
+                assert regional_ip_data is None
+
+        # Start unmodified block from AttnProcessor2_0.
+        # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+        # End of unmodified block from AttnProcessor2_0
+
+        # casting torch.Tensor to torch.FloatTensor to avoid type issues
+        return cast(torch.FloatTensor, hidden_states)

invokeai/backend/stable_diffusion/diffusion/regional_ip_data.py

@@ -0,0 +1,72 @@
+import torch
+
+
+class RegionalIPData:
+    """A class to manage the data for regional IP-Adapter conditioning."""
+
+    def __init__(
+        self,
+        image_prompt_embeds: list[torch.Tensor],
+        scales: list[float],
+        masks: list[torch.Tensor],
+        dtype: torch.dtype,
+        device: torch.device,
+        max_downscale_factor: int = 8,
+    ):
+        """Initialize a `IPAdapterConditioningData` object."""
+        assert len(image_prompt_embeds) == len(scales) == len(masks)
+
+        # The image prompt embeddings.
+        # regional_ip_data[i] contains the image prompt embeddings for the i'th IP-Adapter. Each tensor
+        # has shape (batch_size, num_ip_images, seq_len, ip_embedding_len).
+        self.image_prompt_embeds = image_prompt_embeds
+
+        # The scales for the IP-Adapter attention.
+        # scales[i] contains the attention scale for the i'th IP-Adapter.
+        self.scales = scales
+
+        # The IP-Adapter masks.
+        # self._masks_by_seq_len[s] contains the spatial masks for the downsampling level with query sequence length of
+        # s. It has shape (batch_size, num_ip_images, query_seq_len, 1). The masks have values of 1.0 for included
+        # regions and 0.0 for excluded regions.
+        self._masks_by_seq_len = self._prepare_masks(masks, max_downscale_factor, device, dtype)
+
+    def _prepare_masks(
+        self, masks: list[torch.Tensor], max_downscale_factor: int, device: torch.device, dtype: torch.dtype
+    ) -> dict[int, torch.Tensor]:
+        """Prepare the masks for the IP-Adapter attention."""
+        # Concatenate the masks so that they can be processed more efficiently.
+        mask_tensor = torch.cat(masks, dim=1)
+
+        mask_tensor = mask_tensor.to(device=device, dtype=dtype)
+
+        masks_by_seq_len: dict[int, torch.Tensor] = {}
+
+        # Downsample the spatial dimensions by factors of 2 until max_downscale_factor is reached.
+        downscale_factor = 1
+        while downscale_factor <= max_downscale_factor:
+            b, num_ip_adapters, h, w = mask_tensor.shape
+            # Assert that the batch size is 1, because I haven't thought through batch handling for this feature yet.
+            assert b == 1
+
+            # The IP-Adapters are applied in the cross-attention layers, where the query sequence length is the h * w of
+            # the spatial features.
+            query_seq_len = h * w
+
+            masks_by_seq_len[query_seq_len] = mask_tensor.view((b, num_ip_adapters, -1, 1))
+
+            downscale_factor *= 2
+            if downscale_factor <= max_downscale_factor:
+                # We use max pooling because we downscale to a pretty low resolution, so we don't want small mask
+                # regions to be lost entirely.
+                #
+                # ceil_mode=True is set to mirror the downsampling behavior of SD and SDXL.
+                #
+                # TODO(ryand): In the future, we may want to experiment with other downsampling methods.
+                mask_tensor = torch.nn.functional.max_pool2d(mask_tensor, kernel_size=2, stride=2, ceil_mode=True)
+
+        return masks_by_seq_len
+
+    def get_masks(self, query_seq_len: int) -> torch.Tensor:
+        """Get the mask for the given query sequence length."""
+        return self._masks_by_seq_len[query_seq_len]

invokeai/backend/stable_diffusion/diffusion/regional_prompt_data.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+        TextConditioningRegions,
+    )
+
+
+class RegionalPromptData:
+    """A class to manage the prompt data for regional conditioning."""
+
+    def __init__(
+        self,
+        regions: list[TextConditioningRegions],
+        device: torch.device,
+        dtype: torch.dtype,
+        max_downscale_factor: int = 8,
+    ):
+        """Initialize a `RegionalPromptData` object.
+        Args:
+            regions (list[TextConditioningRegions]): regions[i] contains the prompt regions for the i'th sample in the
+                batch.
+            device (torch.device): The device to use for the attention masks.
+            dtype (torch.dtype): The data type to use for the attention masks.
+            max_downscale_factor: Spatial masks will be prepared for downscale factors from 1 to max_downscale_factor
+                in steps of 2x.
+        """
+        self._regions = regions
+        self._device = device
+        self._dtype = dtype
+        # self._spatial_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query
+        # sequence length of s.
+        self._spatial_masks_by_seq_len: list[dict[int, torch.Tensor]] = self._prepare_spatial_masks(
+            regions, max_downscale_factor
+        )
+        self._negative_cross_attn_mask_score = -10000.0
+
+    def _prepare_spatial_masks(
+        self, regions: list[TextConditioningRegions], max_downscale_factor: int = 8
+    ) -> list[dict[int, torch.Tensor]]:
+        """Prepare the spatial masks for all downscaling factors."""
+        # batch_masks_by_seq_len[b][s] contains the spatial masks for the b'th batch sample with a query sequence length
+        # of s.
+        batch_sample_masks_by_seq_len: list[dict[int, torch.Tensor]] = []
+
+        for batch_sample_regions in regions:
+            batch_sample_masks_by_seq_len.append({})
+
+            batch_sample_masks = batch_sample_regions.masks.to(device=self._device, dtype=self._dtype)
+
+            # Downsample the spatial dimensions by factors of 2 until max_downscale_factor is reached.
+            downscale_factor = 1
+            while downscale_factor <= max_downscale_factor:
+                b, _num_prompts, h, w = batch_sample_masks.shape
+                assert b == 1
+                query_seq_len = h * w
+
+                batch_sample_masks_by_seq_len[-1][query_seq_len] = batch_sample_masks
+
+                downscale_factor *= 2
+                if downscale_factor <= max_downscale_factor:
+                    # We use max pooling because we downscale to a pretty low resolution, so we don't want small prompt
+                    # regions to be lost entirely.
+                    #
+                    # ceil_mode=True is set to mirror the downsampling behavior of SD and SDXL.
+                    #
+                    # TODO(ryand): In the future, we may want to experiment with other downsampling methods (e.g.
+                    # nearest interpolation), and could potentially use a weighted mask rather than a binary mask.
+                    batch_sample_masks = F.max_pool2d(batch_sample_masks, kernel_size=2, stride=2, ceil_mode=True)
+
+        return batch_sample_masks_by_seq_len
+
+    def get_cross_attn_mask(self, query_seq_len: int, key_seq_len: int) -> torch.Tensor:
+        """Get the cross-attention mask for the given query sequence length.
+        Args:
+            query_seq_len: The length of the flattened spatial features at the current downscaling level.
+            key_seq_len (int): The sequence length of the prompt embeddings (which act as the key in the cross-attention
+                layers). This is most likely equal to the max embedding range end, but we pass it explicitly to be sure.
+        Returns:
+            torch.Tensor: The cross-attention score mask.
+                shape: (batch_size, query_seq_len, key_seq_len).
+                dtype: float
+        """
+        batch_size = len(self._spatial_masks_by_seq_len)
+        batch_spatial_masks = [self._spatial_masks_by_seq_len[b][query_seq_len] for b in range(batch_size)]
+
+        # Create an empty attention mask with the correct shape.
+        attn_mask = torch.zeros((batch_size, query_seq_len, key_seq_len), dtype=self._dtype, device=self._device)
+
+        for batch_idx in range(batch_size):
+            batch_sample_spatial_masks = batch_spatial_masks[batch_idx]
+            batch_sample_regions = self._regions[batch_idx]
+
+            # Flatten the spatial dimensions of the mask by reshaping to (1, num_prompts, query_seq_len, 1).
+            _, num_prompts, _, _ = batch_sample_spatial_masks.shape
+            batch_sample_query_masks = batch_sample_spatial_masks.view((1, num_prompts, query_seq_len, 1))
+
+            for prompt_idx, embedding_range in enumerate(batch_sample_regions.ranges):
+                batch_sample_query_scores = batch_sample_query_masks[0, prompt_idx, :, :].clone()
+                batch_sample_query_mask = batch_sample_query_scores > 0.5
+                batch_sample_query_scores[batch_sample_query_mask] = 0.0
+                batch_sample_query_scores[~batch_sample_query_mask] = self._negative_cross_attn_mask_score
+                attn_mask[batch_idx, :, embedding_range.start : embedding_range.end] = batch_sample_query_scores
+
+        return attn_mask

invokeai/backend/stable_diffusion/diffusion/shared_invokeai_diffusion.py

@@ -0,0 +1,496 @@
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, Optional, Union
+
+import torch
+from typing_extensions import TypeAlias
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    IPAdapterData,
+    Range,
+    TextConditioningData,
+    TextConditioningRegions,
+)
+from invokeai.backend.stable_diffusion.diffusion.regional_ip_data import RegionalIPData
+from invokeai.backend.stable_diffusion.diffusion.regional_prompt_data import RegionalPromptData
+
+ModelForwardCallback: TypeAlias = Union[
+    # x, t, conditioning, Optional[cross-attention kwargs]
+    Callable[
+        [torch.Tensor, torch.Tensor, torch.Tensor, Optional[dict[str, Any]]],
+        torch.Tensor,
+    ],
+    Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
+]
+
+
+class InvokeAIDiffuserComponent:
+    """
+    The aim of this component is to provide a single place for code that can be applied identically to
+    all InvokeAI diffusion procedures.
+
+    At the moment it includes the following features:
+    * Cross attention control ("prompt2prompt")
+    * Hybrid conditioning (used for inpainting)
+    """
+
+    debug_thresholding = False
+    sequential_guidance = False
+
+    def __init__(
+        self,
+        model,
+        model_forward_callback: ModelForwardCallback,
+    ):
+        """
+        :param model: the unet model to pass through to cross attention control
+        :param model_forward_callback: a lambda with arguments (x, sigma, conditioning_to_apply). will be called repeatedly. most likely, this should simply call model.forward(x, sigma, conditioning)
+        """
+        config = get_config()
+        self.conditioning = None
+        self.model = model
+        self.model_forward_callback = model_forward_callback
+        self.sequential_guidance = config.sequential_guidance
+
+    def do_controlnet_step(
+        self,
+        control_data,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        step_index: int,
+        total_step_count: int,
+        conditioning_data: TextConditioningData,
+    ):
+        down_block_res_samples, mid_block_res_sample = None, None
+
+        # control_data should be type List[ControlNetData]
+        # this loop covers both ControlNet (one ControlNetData in list)
+        #      and MultiControlNet (multiple ControlNetData in list)
+        for _i, control_datum in enumerate(control_data):
+            control_mode = control_datum.control_mode
+            # soft_injection and cfg_injection are the two ControlNet control_mode booleans
+            #     that are combined at higher level to make control_mode enum
+            #  soft_injection determines whether to do per-layer re-weighting adjustment (if True)
+            #     or default weighting (if False)
+            soft_injection = control_mode == "more_prompt" or control_mode == "more_control"
+            #  cfg_injection = determines whether to apply ControlNet to only the conditional (if True)
+            #      or the default both conditional and unconditional (if False)
+            cfg_injection = control_mode == "more_control" or control_mode == "unbalanced"
+
+            first_control_step = math.floor(control_datum.begin_step_percent * total_step_count)
+            last_control_step = math.ceil(control_datum.end_step_percent * total_step_count)
+            # only apply controlnet if current step is within the controlnet's begin/end step range
+            if step_index >= first_control_step and step_index <= last_control_step:
+                if cfg_injection:
+                    sample_model_input = sample
+                else:
+                    # expand the latents input to control model if doing classifier free guidance
+                    #    (which I think for now is always true, there is conditional elsewhere that stops execution if
+                    #     classifier_free_guidance is <= 1.0 ?)
+                    sample_model_input = torch.cat([sample] * 2)
+
+                added_cond_kwargs = None
+
+                if cfg_injection:  # only applying ControlNet to conditional instead of in unconditioned
+                    if conditioning_data.is_sdxl():
+                        added_cond_kwargs = {
+                            "text_embeds": conditioning_data.cond_text.pooled_embeds,
+                            "time_ids": conditioning_data.cond_text.add_time_ids,
+                        }
+                    encoder_hidden_states = conditioning_data.cond_text.embeds
+                    encoder_attention_mask = None
+                else:
+                    if conditioning_data.is_sdxl():
+                        added_cond_kwargs = {
+                            "text_embeds": torch.cat(
+                                [
+                                    # TODO: how to pad? just by zeros? or even truncate?
+                                    conditioning_data.uncond_text.pooled_embeds,
+                                    conditioning_data.cond_text.pooled_embeds,
+                                ],
+                                dim=0,
+                            ),
+                            "time_ids": torch.cat(
+                                [
+                                    conditioning_data.uncond_text.add_time_ids,
+                                    conditioning_data.cond_text.add_time_ids,
+                                ],
+                                dim=0,
+                            ),
+                        }
+                    (
+                        encoder_hidden_states,
+                        encoder_attention_mask,
+                    ) = self._concat_conditionings_for_batch(
+                        conditioning_data.uncond_text.embeds,
+                        conditioning_data.cond_text.embeds,
+                    )
+                if isinstance(control_datum.weight, list):
+                    # if controlnet has multiple weights, use the weight for the current step
+                    controlnet_weight = control_datum.weight[step_index]
+                else:
+                    # if controlnet has a single weight, use it for all steps
+                    controlnet_weight = control_datum.weight
+
+                # controlnet(s) inference
+                down_samples, mid_sample = control_datum.model(
+                    sample=sample_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=encoder_hidden_states,
+                    controlnet_cond=control_datum.image_tensor,
+                    conditioning_scale=controlnet_weight,  # controlnet specific, NOT the guidance scale
+                    encoder_attention_mask=encoder_attention_mask,
+                    added_cond_kwargs=added_cond_kwargs,
+                    guess_mode=soft_injection,  # this is still called guess_mode in diffusers ControlNetModel
+                    return_dict=False,
+                )
+                if cfg_injection:
+                    # Inferred ControlNet only for the conditional batch.
+                    # To apply the output of ControlNet to both the unconditional and conditional batches,
+                    #    prepend zeros for unconditional batch
+                    down_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_samples]
+                    mid_sample = torch.cat([torch.zeros_like(mid_sample), mid_sample])
+
+                if down_block_res_samples is None and mid_block_res_sample is None:
+                    down_block_res_samples, mid_block_res_sample = down_samples, mid_sample
+                else:
+                    # add controlnet outputs together if have multiple controlnets
+                    down_block_res_samples = [
+                        samples_prev + samples_curr
+                        for samples_prev, samples_curr in zip(down_block_res_samples, down_samples, strict=True)
+                    ]
+                    mid_block_res_sample += mid_sample
+
+        return down_block_res_samples, mid_block_res_sample
+
+    def do_unet_step(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
+        step_index: int,
+        total_step_count: int,
+        down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
+        mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
+        down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
+    ):
+        if self.sequential_guidance:
+            (
+                unconditioned_next_x,
+                conditioned_next_x,
+            ) = self._apply_standard_conditioning_sequentially(
+                x=sample,
+                sigma=timestep,
+                conditioning_data=conditioning_data,
+                ip_adapter_data=ip_adapter_data,
+                step_index=step_index,
+                total_step_count=total_step_count,
+                down_block_additional_residuals=down_block_additional_residuals,
+                mid_block_additional_residual=mid_block_additional_residual,
+                down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+            )
+        else:
+            (
+                unconditioned_next_x,
+                conditioned_next_x,
+            ) = self._apply_standard_conditioning(
+                x=sample,
+                sigma=timestep,
+                conditioning_data=conditioning_data,
+                ip_adapter_data=ip_adapter_data,
+                step_index=step_index,
+                total_step_count=total_step_count,
+                down_block_additional_residuals=down_block_additional_residuals,
+                mid_block_additional_residual=mid_block_additional_residual,
+                down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+            )
+
+        return unconditioned_next_x, conditioned_next_x
+
+    def _concat_conditionings_for_batch(self, unconditioning, conditioning):
+        def _pad_conditioning(cond, target_len, encoder_attention_mask):
+            conditioning_attention_mask = torch.ones(
+                (cond.shape[0], cond.shape[1]), device=cond.device, dtype=cond.dtype
+            )
+
+            if cond.shape[1] < max_len:
+                conditioning_attention_mask = torch.cat(
+                    [
+                        conditioning_attention_mask,
+                        torch.zeros((cond.shape[0], max_len - cond.shape[1]), device=cond.device, dtype=cond.dtype),
+                    ],
+                    dim=1,
+                )
+
+                cond = torch.cat(
+                    [
+                        cond,
+                        torch.zeros(
+                            (cond.shape[0], max_len - cond.shape[1], cond.shape[2]),
+                            device=cond.device,
+                            dtype=cond.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+
+            if encoder_attention_mask is None:
+                encoder_attention_mask = conditioning_attention_mask
+            else:
+                encoder_attention_mask = torch.cat(
+                    [
+                        encoder_attention_mask,
+                        conditioning_attention_mask,
+                    ]
+                )
+
+            return cond, encoder_attention_mask
+
+        encoder_attention_mask = None
+        if unconditioning.shape[1] != conditioning.shape[1]:
+            max_len = max(unconditioning.shape[1], conditioning.shape[1])
+            unconditioning, encoder_attention_mask = _pad_conditioning(unconditioning, max_len, encoder_attention_mask)
+            conditioning, encoder_attention_mask = _pad_conditioning(conditioning, max_len, encoder_attention_mask)
+
+        return torch.cat([unconditioning, conditioning]), encoder_attention_mask
+
+    # methods below are called from do_diffusion_step and should be considered private to this class.
+
+    def _apply_standard_conditioning(
+        self,
+        x: torch.Tensor,
+        sigma: torch.Tensor,
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
+        step_index: int,
+        total_step_count: int,
+        down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
+        mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
+        down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Runs the conditioned and unconditioned UNet forward passes in a single batch for faster inference speed at
+        the cost of higher memory usage.
+        """
+        x_twice = torch.cat([x] * 2)
+        sigma_twice = torch.cat([sigma] * 2)
+
+        cross_attention_kwargs = {}
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+            # Note that we 'stack' to produce tensors of shape (batch_size, num_ip_images, seq_len, token_len).
+            image_prompt_embeds = [
+                torch.stack([ipa_conditioning.uncond_image_prompt_embeds, ipa_conditioning.cond_image_prompt_embeds])
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
+
+        added_cond_kwargs = None
+        if conditioning_data.is_sdxl():
+            added_cond_kwargs = {
+                "text_embeds": torch.cat(
+                    [
+                        # TODO: how to pad? just by zeros? or even truncate?
+                        conditioning_data.uncond_text.pooled_embeds,
+                        conditioning_data.cond_text.pooled_embeds,
+                    ],
+                    dim=0,
+                ),
+                "time_ids": torch.cat(
+                    [
+                        conditioning_data.uncond_text.add_time_ids,
+                        conditioning_data.cond_text.add_time_ids,
+                    ],
+                    dim=0,
+                ),
+            }
+
+        if conditioning_data.cond_regions is not None or conditioning_data.uncond_regions is not None:
+            # TODO(ryand): We currently initialize RegionalPromptData for every denoising step. The text conditionings
+            # and masks are not changing from step-to-step, so this really only needs to be done once. While this seems
+            # painfully inefficient, the time spent is typically negligible compared to the forward inference pass of
+            # the UNet. The main reason that this hasn't been moved up to eliminate redundancy is that it is slightly
+            # awkward to handle both standard conditioning and sequential conditioning further up the stack.
+            regions = []
+            for c, r in [
+                (conditioning_data.uncond_text, conditioning_data.uncond_regions),
+                (conditioning_data.cond_text, conditioning_data.cond_regions),
+            ]:
+                if r is None:
+                    # Create a dummy mask and range for text conditioning that doesn't have region masks.
+                    _, _, h, w = x.shape
+                    r = TextConditioningRegions(
+                        masks=torch.ones((1, 1, h, w), dtype=x.dtype),
+                        ranges=[Range(start=0, end=c.embeds.shape[1])],
+                    )
+                regions.append(r)
+
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=regions, device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
+        both_conditionings, encoder_attention_mask = self._concat_conditionings_for_batch(
+            conditioning_data.uncond_text.embeds, conditioning_data.cond_text.embeds
+        )
+        both_results = self.model_forward_callback(
+            x_twice,
+            sigma_twice,
+            both_conditionings,
+            cross_attention_kwargs=cross_attention_kwargs,
+            encoder_attention_mask=encoder_attention_mask,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+            down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+            added_cond_kwargs=added_cond_kwargs,
+        )
+        unconditioned_next_x, conditioned_next_x = both_results.chunk(2)
+        return unconditioned_next_x, conditioned_next_x
+
+    def _apply_standard_conditioning_sequentially(
+        self,
+        x: torch.Tensor,
+        sigma,
+        conditioning_data: TextConditioningData,
+        ip_adapter_data: Optional[list[IPAdapterData]],
+        step_index: int,
+        total_step_count: int,
+        down_block_additional_residuals: Optional[torch.Tensor] = None,  # for ControlNet
+        mid_block_additional_residual: Optional[torch.Tensor] = None,  # for ControlNet
+        down_intrablock_additional_residuals: Optional[torch.Tensor] = None,  # for T2I-Adapter
+    ):
+        """Runs the conditioned and unconditioned UNet forward passes sequentially for lower memory usage at the cost of
+        slower execution speed.
+        """
+        # Since we are running the conditioned and unconditioned passes sequentially, we need to split the ControlNet
+        # and T2I-Adapter residuals into two chunks.
+        uncond_down_block, cond_down_block = None, None
+        if down_block_additional_residuals is not None:
+            uncond_down_block, cond_down_block = [], []
+            for down_block in down_block_additional_residuals:
+                _uncond_down, _cond_down = down_block.chunk(2)
+                uncond_down_block.append(_uncond_down)
+                cond_down_block.append(_cond_down)
+
+        uncond_down_intrablock, cond_down_intrablock = None, None
+        if down_intrablock_additional_residuals is not None:
+            uncond_down_intrablock, cond_down_intrablock = [], []
+            for down_intrablock in down_intrablock_additional_residuals:
+                _uncond_down, _cond_down = down_intrablock.chunk(2)
+                uncond_down_intrablock.append(_uncond_down)
+                cond_down_intrablock.append(_cond_down)
+
+        uncond_mid_block, cond_mid_block = None, None
+        if mid_block_additional_residual is not None:
+            uncond_mid_block, cond_mid_block = mid_block_additional_residual.chunk(2)
+
+        #####################
+        # Unconditioned pass
+        #####################
+
+        cross_attention_kwargs = {}
+
+        # Prepare IP-Adapter cross-attention kwargs for the unconditioned pass.
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+            # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
+            image_prompt_embeds = [
+                torch.unsqueeze(ipa_conditioning.uncond_image_prompt_embeds, dim=0)
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
+
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
+
+        # Prepare SDXL conditioning kwargs for the unconditioned pass.
+        added_cond_kwargs = None
+        if conditioning_data.is_sdxl():
+            added_cond_kwargs = {
+                "text_embeds": conditioning_data.uncond_text.pooled_embeds,
+                "time_ids": conditioning_data.uncond_text.add_time_ids,
+            }
+
+        # Prepare prompt regions for the unconditioned pass.
+        if conditioning_data.uncond_regions is not None:
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=[conditioning_data.uncond_regions], device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
+        # Run unconditioned UNet denoising (i.e. negative prompt).
+        unconditioned_next_x = self.model_forward_callback(
+            x,
+            sigma,
+            conditioning_data.uncond_text.embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=uncond_down_block,
+            mid_block_additional_residual=uncond_mid_block,
+            down_intrablock_additional_residuals=uncond_down_intrablock,
+            added_cond_kwargs=added_cond_kwargs,
+        )
+
+        ###################
+        # Conditioned pass
+        ###################
+
+        cross_attention_kwargs = {}
+
+        if ip_adapter_data is not None:
+            ip_adapter_conditioning = [ipa.ip_adapter_conditioning for ipa in ip_adapter_data]
+            # Note that we 'unsqueeze' to produce tensors of shape (batch_size=1, num_ip_images, seq_len, token_len).
+            image_prompt_embeds = [
+                torch.unsqueeze(ipa_conditioning.cond_image_prompt_embeds, dim=0)
+                for ipa_conditioning in ip_adapter_conditioning
+            ]
+
+            scales = [ipa.scale_for_step(step_index, total_step_count) for ipa in ip_adapter_data]
+            ip_masks = [ipa.mask for ipa in ip_adapter_data]
+            regional_ip_data = RegionalIPData(
+                image_prompt_embeds=image_prompt_embeds, scales=scales, masks=ip_masks, dtype=x.dtype, device=x.device
+            )
+            cross_attention_kwargs["regional_ip_data"] = regional_ip_data
+
+        # Prepare SDXL conditioning kwargs for the conditioned pass.
+        added_cond_kwargs = None
+        if conditioning_data.is_sdxl():
+            added_cond_kwargs = {
+                "text_embeds": conditioning_data.cond_text.pooled_embeds,
+                "time_ids": conditioning_data.cond_text.add_time_ids,
+            }
+
+        # Prepare prompt regions for the conditioned pass.
+        if conditioning_data.cond_regions is not None:
+            cross_attention_kwargs["regional_prompt_data"] = RegionalPromptData(
+                regions=[conditioning_data.cond_regions], device=x.device, dtype=x.dtype
+            )
+            cross_attention_kwargs["percent_through"] = step_index / total_step_count
+
+        # Run conditioned UNet denoising (i.e. positive prompt).
+        conditioned_next_x = self.model_forward_callback(
+            x,
+            sigma,
+            conditioning_data.cond_text.embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=cond_down_block,
+            mid_block_additional_residual=cond_mid_block,
+            down_intrablock_additional_residuals=cond_down_intrablock,
+            added_cond_kwargs=added_cond_kwargs,
+        )
+        return unconditioned_next_x, conditioned_next_x
+
+    def _combine(self, unconditioned_next_x, conditioned_next_x, guidance_scale):
+        # to scale how much effect conditioning has, calculate the changes it does and then scale that
+        scaled_delta = (conditioned_next_x - unconditioned_next_x) * guidance_scale
+        combined_next_x = unconditioned_next_x + scaled_delta
+        return combined_next_x

invokeai/backend/stable_diffusion/diffusion/unet_attention_patcher.py

@@ -0,0 +1,75 @@
+from contextlib import contextmanager
+from typing import List, Optional, TypedDict
+
+from diffusers.models import UNet2DConditionModel
+
+from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+from invokeai.backend.stable_diffusion.diffusion.custom_atttention import (
+    CustomAttnProcessor2_0,
+    IPAdapterAttentionWeights,
+)
+
+
+class UNetIPAdapterData(TypedDict):
+    ip_adapter: IPAdapter
+    target_blocks: List[str]  # Blocks where IP-Adapter should be applied
+    method: str  # Style or other method type
+
+
+class UNetAttentionPatcher:
+    """A class for patching a UNet with CustomAttnProcessor2_0 attention layers."""
+
+    def __init__(self, ip_adapter_data: Optional[List[UNetIPAdapterData]]):
+        self._ip_adapters = ip_adapter_data
+
+    def _prepare_attention_processors(self, unet: UNet2DConditionModel):
+        """Prepare a dict of attention processors that can be injected into a unet, and load the IP-Adapter attention
+        weights into them (if IP-Adapters are being applied).
+        Note that the `unet` param is only used to determine attention block dimensions and naming.
+        """
+        # Construct a dict of attention processors based on the UNet's architecture.
+        attn_procs = {}
+        for idx, name in enumerate(unet.attn_processors.keys()):
+            if name.endswith("attn1.processor") or self._ip_adapters is None:
+                # "attn1" processors do not use IP-Adapters.
+                attn_procs[name] = CustomAttnProcessor2_0()
+            else:
+                # Collect the weights from each IP Adapter for the idx'th attention processor.
+                ip_adapter_attention_weights_collection: list[IPAdapterAttentionWeights] = []
+
+                for ip_adapter in self._ip_adapters:
+                    ip_adapter_weights = ip_adapter["ip_adapter"].attn_weights.get_attention_processor_weights(idx)
+                    skip = True
+                    negative = False
+                    for block in ip_adapter["target_blocks"]:
+                        if block in name:
+                            skip = False
+                            negative = ip_adapter["method"] == "style_precise" and (
+                                block == "down_blocks.2.attentions.1"
+                                or block == "down_blocks.2"
+                                or block == "mid_block"
+                            )
+                            break
+                    ip_adapter_attention_weights: IPAdapterAttentionWeights = IPAdapterAttentionWeights(
+                        ip_adapter_weights=ip_adapter_weights, skip=skip, negative=negative
+                    )
+                    ip_adapter_attention_weights_collection.append(ip_adapter_attention_weights)
+
+                attn_procs[name] = CustomAttnProcessor2_0(ip_adapter_attention_weights_collection)
+
+        return attn_procs
+
+    @contextmanager
+    def apply_ip_adapter_attention(self, unet: UNet2DConditionModel):
+        """A context manager that patches `unet` with CustomAttnProcessor2_0 attention layers."""
+        attn_procs = self._prepare_attention_processors(unet)
+        orig_attn_processors = unet.attn_processors
+
+        try:
+            # Note to future devs: set_attn_processor(...) does something slightly unexpected - it pops elements from
+            # the passed dict. So, if you wanted to keep the dict for future use, you'd have to make a
+            # moderately-shallow copy of it. E.g. `attn_procs_copy = {k: v for k, v in attn_procs.items()}`.
+            unet.set_attn_processor(attn_procs)
+            yield None
+        finally:
+            unet.set_attn_processor(orig_attn_processors)