Create a UNetAttentionPatcher for patching UNet models with CustomAttnProcessor2_0 modules.

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -22,12 +22,12 @@ from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.ip_adapter.unet_patcher import UNetPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     IPAdapterConditioningInfo,
     TextConditioningData,
 )
 from invokeai.backend.stable_diffusion.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
+from invokeai.backend.stable_diffusion.diffusion.unet_attention_patcher import UNetAttentionPatcher
 from invokeai.backend.util.attention import auto_detect_slice_size
 from invokeai.backend.util.devices import normalize_device
 
@@ -412,7 +412,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         elif ip_adapter_data is not None:
             # TODO(ryand): Should we raise an exception if both custom attention and IP-Adapter attention are active?
             # As it is now, the IP-Adapter will silently be skipped.
-            ip_adapter_unet_patcher = UNetPatcher([ipa.ip_adapter_model for ipa in ip_adapter_data])
+            ip_adapter_unet_patcher = UNetAttentionPatcher([ipa.ip_adapter_model for ipa in ip_adapter_data])
             attn_ctx = ip_adapter_unet_patcher.apply_ip_adapter_attention(self.invokeai_diffuser.model)
             self.use_ip_adapter = True
         else:
@@ -476,7 +476,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         control_data: List[ControlNetData] = None,
         ip_adapter_data: Optional[list[IPAdapterData]] = None,
         t2i_adapter_data: Optional[list[T2IAdapterData]] = None,
-        ip_adapter_unet_patcher: Optional[UNetPatcher] = None,
+        ip_adapter_unet_patcher: Optional[UNetAttentionPatcher] = None,
     ):
         # invokeai_diffuser has batched timesteps, but diffusers schedulers expect a single value
         timestep = t[0]

invokeai/backend/stable_diffusion/diffusion/unet_attention_patcher.py

@@ -0,0 +1,55 @@
+from contextlib import contextmanager
+from typing import Optional
+
+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
+
+
+class UNetAttentionPatcher:
+    """A class for patching a UNet with CustomAttnProcessor2_0 attention layers."""
+
+    def __init__(self, ip_adapters: Optional[list[IPAdapter]]):
+        self._ip_adapters = ip_adapters
+        self._ip_adapter_scales = None
+
+        if self._ip_adapters is not None:
+            self._ip_adapter_scales = [1.0] * len(self._ip_adapters)
+
+    def set_scale(self, idx: int, value: float):
+        self._ip_adapter_scales[idx] = value
+
+    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.
+                attn_procs[name] = CustomAttnProcessor2_0(
+                    [ip_adapter.attn_weights.get_attention_processor_weights(idx) for ip_adapter in self._ip_adapters],
+                    self._ip_adapter_scales,
+                )
+        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)