hacked version of attention_processor.py that batches ip_adapters

proof of concept implementation of hacking the
attention processor. when it detects a batch in
which the batch_size matches the number of ip_adapters,
it attempts to run each ip_adapter on one element
of the batch instead of squashing all ip_adapters together.

Author: dribnet

src/diffusers/models/attention_processor.py

@@ -3955,68 +3955,106 @@ def __call__(
         else:
             ip_adapter_masks = [None] * len(self.scale)
 
-        # for ip-adapter
-        for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(
-            ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks
-        ):
-            skip = False
-            if isinstance(scale, list):
-                if all(s == 0 for s in scale):
-                    skip = True
-            elif scale == 0:
-                skip = True
-            if not skip:
-                if mask is not None:
-                    if not isinstance(scale, list):
-                        scale = [scale] * mask.shape[1]
+        # HACK: if the number of ip_adapters match the batch, process as a special case
+        num_hidden_states = len(ip_hidden_states)
+        if batch_size == 2 * num_hidden_states:
+            result_keys = []
+            result_values = []
+            for i in range(batch_size):
+                current_ip_hidden_states = ip_hidden_states[i % num_hidden_states]
+                scale = self.scale[i % num_hidden_states]
+                to_k_ip = self.to_k_ip[i % num_hidden_states]
+                to_v_ip = self.to_v_ip[i % num_hidden_states]
+                mask = ip_adapter_masks[i % num_hidden_states]
+
+                ip_key = to_k_ip(current_ip_hidden_states[i, :, :, :])
+                ip_value = to_v_ip(current_ip_hidden_states[i, :, :, :])
+
+                ip_key = ip_key.view(2, -1, attn.heads, head_dim).transpose(1, 2)
+                ip_value = ip_value.view(2, -1, attn.heads, head_dim).transpose(1, 2)
+
+                result_keys.append(ip_key[0])
+                result_values.append(ip_value[0])
+
+            ip_key = torch.stack(result_keys, dim=0)
+            ip_value = torch.stack(result_values, dim=0)
+
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            # TODO: add support for attn.scale when we move to Torch 2.1
+            current_ip_hidden_states = F.scaled_dot_product_attention(
+                query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+            )
 
-                    current_num_images = mask.shape[1]
-                    for i in range(current_num_images):
-                        ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])
-                        ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])
+            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(query.dtype)
+
+            hidden_states = hidden_states + scale * current_ip_hidden_states
+
+        else:
+            # for ip-adapter
+            for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(
+                ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks
+            ):
+                skip = False
+                if isinstance(scale, list):
+                    if all(s == 0 for s in scale):
+                        skip = True
+                elif scale == 0:
+                    skip = True
+                if not skip:
+                    if mask is not None:
+                        if not isinstance(scale, list):
+                            scale = [scale] * mask.shape[1]
+
+                        current_num_images = mask.shape[1]
+                        for i in range(current_num_images):
+                            ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])
+                            ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])
+
+                            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)
+
+                            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+                            # TODO: add support for attn.scale when we move to Torch 2.1
+                            _current_ip_hidden_states = F.scaled_dot_product_attention(
+                                query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+                            )
+
+                            _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(query.dtype)
+
+                            mask_downsample = IPAdapterMaskProcessor.downsample(
+                                mask[:, i, :, :],
+                                batch_size,
+                                _current_ip_hidden_states.shape[1],
+                                _current_ip_hidden_states.shape[2],
+                            )
+
+                            mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)
+                            hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)
+                    else:
+                        ip_key = to_k_ip(current_ip_hidden_states)
+                        ip_value = to_v_ip(current_ip_hidden_states)
 
                         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)
 
                         # the output of sdp = (batch, num_heads, seq_len, head_dim)
                         # TODO: add support for attn.scale when we move to Torch 2.1
-                        _current_ip_hidden_states = F.scaled_dot_product_attention(
+                        current_ip_hidden_states = F.scaled_dot_product_attention(
                             query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
                         )
 
-                        _current_ip_hidden_states = _current_ip_hidden_states.transpose(1, 2).reshape(
+                        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(query.dtype)
-
-                        mask_downsample = IPAdapterMaskProcessor.downsample(
-                            mask[:, i, :, :],
-                            batch_size,
-                            _current_ip_hidden_states.shape[1],
-                            _current_ip_hidden_states.shape[2],
-                        )
-
-                        mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)
-                        hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)
-                else:
-                    ip_key = to_k_ip(current_ip_hidden_states)
-                    ip_value = to_v_ip(current_ip_hidden_states)
-
-                    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)
+                        current_ip_hidden_states = current_ip_hidden_states.to(query.dtype)
 
-                    # the output of sdp = (batch, num_heads, seq_len, head_dim)
-                    # TODO: add support for attn.scale when we move to Torch 2.1
-                    current_ip_hidden_states = F.scaled_dot_product_attention(
-                        query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
-                    )
-
-                    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(query.dtype)
-
-                    hidden_states = hidden_states + scale * current_ip_hidden_states
+                        hidden_states = hidden_states + scale * current_ip_hidden_states
 
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)