Initial implementation of Flux multi IP-Adapter

Author: guiyrt

examples/community/pipeline_flux_semantic_guidance.py

@@ -537,9 +537,9 @@ def prepare_ip_adapter_image_embeds(
             if not isinstance(ip_adapter_image, list):
                 ip_adapter_image = [ip_adapter_image]
 
-            if len(ip_adapter_image) != len(self.transformer.encoder_hid_proj.image_projection_layers):
+            if len(ip_adapter_image) != self.transformer.encoder_hid_proj.num_ip_adapters:
                 raise ValueError(
-                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.transformer.encoder_hid_proj.image_projection_layers)} IP Adapters."
+                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters."
                 )
 
             for single_ip_adapter_image, image_proj_layer in zip(

src/diffusers/loaders/ip_adapter.py

@@ -577,29 +577,38 @@ def LinearStrengthModel(start, finish, size):
         pipeline.set_ip_adapter_scale(ip_strengths)
         ```
         """
-        transformer = self.transformer
-        if not isinstance(scale, list):
-            scale = [[scale] * transformer.config.num_layers]
-        elif isinstance(scale, list) and isinstance(scale[0], int) or isinstance(scale[0], float):
-            if len(scale) != transformer.config.num_layers:
-                raise ValueError(f"Expected list of {transformer.config.num_layers} scales, got {len(scale)}.")
+
+        from ..pipelines.pipeline_loading_utils import _get_detailed_type, _is_valid_type
+
+        scale_type = Union[int, float]
+        num_ip_adapters = self.transformer.encoder_hid_proj.num_ip_adapters
+        num_layers = self.transformer.config.num_layers
+
+        # Single value for all layers of all IP-Adapters
+        if isinstance(scale, scale_type):
+            scale = [scale for _ in range(num_ip_adapters)]
+        # List of per-layer scales for a single IP-Adapter
+        elif _is_valid_type(scale, List[scale_type]) and num_ip_adapters == 1:
             scale = [scale]
+        # Invalid scale type
+        elif not _is_valid_type(scale, List[Union[scale_type, List[scale_type]]]):
+            raise TypeError(f"Unexpected type {_get_detailed_type(scale)} for scale.")
 
-        scale_configs = scale
+        if len(scale) != num_ip_adapters:
+            raise ValueError(f"Cannot assign {len(scale)} scales to {num_ip_adapters} IP-Adapters.")
 
-        key_id = 0
-        for attn_name, attn_processor in transformer.attn_processors.items():
-            if isinstance(attn_processor, (FluxIPAdapterJointAttnProcessor2_0)):
-                if len(scale_configs) != len(attn_processor.scale):
-                    raise ValueError(
-                        f"Cannot assign {len(scale_configs)} scale_configs to "
-                        f"{len(attn_processor.scale)} IP-Adapter."
-                    )
-                elif len(scale_configs) == 1:
-                    scale_configs = scale_configs * len(attn_processor.scale)
-                for i, scale_config in enumerate(scale_configs):
-                    attn_processor.scale[i] = scale_config[key_id]
-                key_id += 1
+        if any(len(s) != num_layers for s in scale if isinstance(s, list)):
+            invalid_scale_sizes = {len(s) for s in scale if isinstance(s, list)} - {num_layers}
+            raise ValueError(
+                f"Expected list of {num_layers} scales, got {', '.join(str(x) for x in invalid_scale_sizes)}."
+            )
+
+        # Scalars are transformed to lists with length num_layers
+        scale_configs = [[s] * num_layers if isinstance(s, scale_type) else s for s in scale]
+
+        # Set scales. zip over scale_configs prevents going into single transformer layers
+        for attn_processor, *scale in zip(self.transformer.attn_processors.values(), *scale_configs):
+            attn_processor.scale = scale
 
     def unload_ip_adapter(self):
         """

src/diffusers/models/attention_processor.py

@@ -2778,9 +2778,8 @@ def __call__(
 
             # IP-adapter
             ip_query = hidden_states_query_proj
-            ip_attn_output = None
-            # for ip-adapter
-            # TODO: support for multiple adapters
+            ip_attn_output = torch.zeros_like(hidden_states)
+
             for current_ip_hidden_states, scale, to_k_ip, to_v_ip in zip(
                 ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip
             ):
@@ -2791,12 +2790,14 @@ def __call__(
                 ip_value = ip_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
-                ip_attn_output = F.scaled_dot_product_attention(
+                current_ip_hidden_states = F.scaled_dot_product_attention(
                     ip_query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
                 )
-                ip_attn_output = ip_attn_output.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-                ip_attn_output = scale * ip_attn_output
-                ip_attn_output = ip_attn_output.to(ip_query.dtype)
+                current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape(
+                    batch_size, -1, attn.heads * head_dim
+                )
+                current_ip_hidden_states = current_ip_hidden_states.to(ip_query.dtype)
+                ip_attn_output += scale * current_ip_hidden_states
 
             return hidden_states, encoder_hidden_states, ip_attn_output
         else:

src/diffusers/models/embeddings.py

@@ -2583,6 +2583,11 @@ def __init__(self, IPAdapterImageProjectionLayers: Union[List[nn.Module], Tuple[
         super().__init__()
         self.image_projection_layers = nn.ModuleList(IPAdapterImageProjectionLayers)
 
+    @property
+    def num_ip_adapters(self) -> int:
+        """Number of IP-Adapters loaded."""
+        return len(self.image_projection_layers)
+
     def forward(self, image_embeds: List[torch.Tensor]):
         projected_image_embeds = []
 

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -405,9 +405,9 @@ def prepare_ip_adapter_image_embeds(
             if not isinstance(ip_adapter_image, list):
                 ip_adapter_image = [ip_adapter_image]
 
-            if len(ip_adapter_image) != len(self.transformer.encoder_hid_proj.image_projection_layers):
+            if len(ip_adapter_image) != self.transformer.encoder_hid_proj.num_ip_adapters:
                 raise ValueError(
-                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.transformer.encoder_hid_proj.image_projection_layers)} IP Adapters."
+                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters."
                 )
 
             for single_ip_adapter_image, image_proj_layer in zip(
@@ -868,14 +868,23 @@ def __call__(
         else:
             guidance = None
 
+        # TODO: Clarify this section
         if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
             negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
         ):
-            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            negative_ip_adapter_image = (
+                [np.zeros((width, height, 3), dtype=np.uint8) for _ in range(len(ip_adapter_image))]
+                if isinstance(ip_adapter_image, list)
+                else np.zeros((width, height, 3), dtype=np.uint8)
+            )
         elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
             negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
         ):
-            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            ip_adapter_image = (
+                [np.zeros((width, height, 3), dtype=np.uint8) for _ in range(len(negative_ip_adapter_image))]
+                if isinstance(negative_ip_adapter_image, list)
+                else np.zeros((width, height, 3), dtype=np.uint8)
+            )
 
         if self.joint_attention_kwargs is None:
             self._joint_attention_kwargs = {}

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -17,7 +17,7 @@
 import re
 import warnings
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union, get_args, get_origin
 
 import requests
 import torch
@@ -1059,3 +1059,76 @@ def _maybe_raise_error_for_incorrect_transformers(config_dict):
                 break
     if has_transformers_component and not is_transformers_version(">", "4.47.1"):
         raise ValueError("Please upgrade your `transformers` installation to the latest version to use DDUF.")
+
+
+def _is_valid_type(obj: Any, class_or_tuple: Union[Type, Tuple[Type, ...]]) -> bool:
+    """
+    Checks if an object is an instance of any of the provided types. For collections, it checks if every element is of
+    the correct type as well.
+    """
+    if not isinstance(class_or_tuple, tuple):
+        class_or_tuple = (class_or_tuple,)
+
+    # Unpack unions
+    unpacked_class_or_tuple = []
+    for t in class_or_tuple:
+        if get_origin(t) is Union:
+            unpacked_class_or_tuple.extend(get_args(t))
+        else:
+            unpacked_class_or_tuple.append(t)
+    class_or_tuple = tuple(unpacked_class_or_tuple)
+
+    if Any in class_or_tuple:
+        return True
+
+    obj_type = type(obj)
+    # Classes with obj's type
+    class_or_tuple = {t for t in class_or_tuple if isinstance(obj, get_origin(t) or t)}
+
+    # Singular types (e.g. int, ControlNet, ...)
+    # Untyped collections (e.g. List, but not List[int])
+    elem_class_or_tuple = {get_args(t) for t in class_or_tuple}
+    if () in elem_class_or_tuple:
+        return True
+    # Typed lists or sets
+    elif obj_type in (list, set):
+        return any(all(_is_valid_type(x, t) for x in obj) for t in elem_class_or_tuple)
+    # Typed tuples
+    elif obj_type is tuple:
+        return any(
+            # Tuples with any length and single type (e.g. Tuple[int, ...])
+            (len(t) == 2 and t[-1] is Ellipsis and all(_is_valid_type(x, t[0]) for x in obj))
+            or
+            # Tuples with fixed length and any types (e.g. Tuple[int, str])
+            (len(obj) == len(t) and all(_is_valid_type(x, tt) for x, tt in zip(obj, t)))
+            for t in elem_class_or_tuple
+        )
+    # Typed dicts
+    elif obj_type is dict:
+        return any(
+            all(_is_valid_type(k, kt) and _is_valid_type(v, vt) for k, v in obj.items())
+            for kt, vt in elem_class_or_tuple
+        )
+
+    else:
+        return False
+
+
+def _get_detailed_type(obj: Any) -> Type:
+    """
+    Gets a detailed type for an object, including nested types for collections.
+    """
+    obj_type = type(obj)
+
+    if obj_type in (list, set):
+        obj_origin_type = List if obj_type is list else Set
+        elems_type = Union[tuple({_get_detailed_type(x) for x in obj})]
+        return obj_origin_type[elems_type]
+    elif obj_type is tuple:
+        return Tuple[tuple(_get_detailed_type(x) for x in obj)]
+    elif obj_type is dict:
+        keys_type = Union[tuple({_get_detailed_type(k) for k in obj.keys()})]
+        values_type = Union[tuple({_get_detailed_type(k) for k in obj.values()})]
+        return Dict[keys_type, values_type]
+    else:
+        return obj_type