test

Author: hlky

src/diffusers/models/attention_processor.py

@@ -2604,15 +2604,15 @@ def __call__(
             ip_hidden_states = image_projection
 
             ip_query = hidden_states_query_proj
-            ip_attn_output = None
+            ip_attn_outputs = []
             # for ip-adapter
             # TODO: fix for multiple
-            # NOTE: run zeros image embed at the same time?
             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
             ):
-                ip_key = to_k_ip(current_ip_hidden_states)
-                ip_value = to_v_ip(current_ip_hidden_states)
+                positive_ip, negative_ip = current_ip_hidden_states
+                ip_key = to_k_ip(positive_ip)
+                ip_value = to_v_ip(positive_ip)
 
                 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)
@@ -2624,8 +2624,24 @@ def __call__(
                 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)
+                ip_attn_outputs.append(ip_attn_output)
 
-            return hidden_states, encoder_hidden_states, ip_attn_output
+                ip_key = to_k_ip(negative_ip)
+                ip_value = to_v_ip(negative_ip)
+
+                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
+                ip_attn_output = 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)
+                ip_attn_outputs.append(ip_attn_output)
+
+            return hidden_states, encoder_hidden_states, ip_attn_outputs
         else:
             return hidden_states
 

src/diffusers/models/transformers/transformer_flux.py

@@ -187,7 +187,7 @@ def forward(
         if len(attention_outputs) == 2:
             attn_output, context_attn_output = attention_outputs
         elif len(attention_outputs) == 3:
-            attn_output, context_attn_output, ip_attn_output = attention_outputs
+            attn_output, context_attn_output, ip_attn_outputs = attention_outputs
 
         # Process attention outputs for the `hidden_states`.
         attn_output = gate_msa.unsqueeze(1) * attn_output
@@ -201,7 +201,8 @@ def forward(
 
         hidden_states = hidden_states + ff_output
         if len(attention_outputs) == 3:
-            hidden_states = hidden_states + ip_attn_output
+            positive_ip_attn, negative_ip_attn = ip_attn_outputs
+            hidden_states = hidden_states + positive_ip_attn + negative_ip_attn
 
         # Process attention outputs for the `encoder_hidden_states`.
 

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -401,9 +401,11 @@ def encode_image(self, image, device, num_images_per_prompt):
         return image_embeds
 
     def prepare_ip_adapter_image_embeds(
-        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt
+        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, height, width, dtype
     ):
         image_embeds = []
+        negative_embeds = []
+        negative_image = np.zeros((width, height, 3), dtype=np.uint8)
         if ip_adapter_image_embeds is None:
             if not isinstance(ip_adapter_image, list):
                 ip_adapter_image = [ip_adapter_image]
@@ -417,19 +419,27 @@ def prepare_ip_adapter_image_embeds(
                 ip_adapter_image, self.transformer.encoder_hid_proj.image_projection_layers
             ):
                 single_image_embeds = self.encode_image(single_ip_adapter_image, device, 1)
+                negative_image_embeds = self.encode_image(negative_image, device, 1)
 
                 image_embeds.append(single_image_embeds[None, :])
                 image_embeds = self.transformer.encoder_hid_proj(image_embeds)
+                negative_embeds.append(negative_image_embeds[None, :])
+                negative_embeds = self.transformer.encoder_hid_proj(negative_embeds)
         else:
             for single_image_embeds in ip_adapter_image_embeds:
                 image_embeds = self.transformer.encoder_hid_proj(single_image_embeds)
                 image_embeds.append(single_image_embeds)
+                negative_image_embeds = self.encode_image(negative_image, device, 1)
+                negative_embeds.append(negative_image_embeds[None, :])
+                negative_embeds = self.transformer.encoder_hid_proj(negative_embeds)
 
         ip_adapter_image_embeds = []
-        for i, single_image_embeds in enumerate(image_embeds):
+        for i, (single_image_embeds, negative_image_embed) in enumerate(zip(image_embeds, negative_embeds)):
             single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
-            single_image_embeds = single_image_embeds.to(device=device)
-            ip_adapter_image_embeds.append(single_image_embeds)
+            single_image_embeds = single_image_embeds.to(device=device, dtype=dtype)
+            negative_image_embed = torch.cat([negative_image_embed] * num_images_per_prompt, dim=0)
+            negative_image_embed = negative_image_embed.to(device=device, dtype=dtype)
+            ip_adapter_image_embeds.append((single_image_embeds, negative_image_embed))
 
         return ip_adapter_image_embeds
 
@@ -794,6 +804,9 @@ def __call__(
                 ip_adapter_image_embeds,
                 device,
                 batch_size * num_images_per_prompt,
+                height,
+                width,
+                latents.dtype,
             )
             if self.joint_attention_kwargs is None:
                 self._joint_attention_kwargs = {}