cross-attention mask

Author: Birch-san

src/diffusers/models/attention.py

@@ -181,6 +181,7 @@ def forward(
         hidden_states,
         encoder_hidden_states=None,
         timestep=None,
+        cross_attn_mask: Optional[torch.Tensor] = None,
         cross_attention_kwargs=None,
         return_dict: bool = True,
     ):
@@ -225,6 +226,7 @@ def forward(
                 hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
                 timestep=timestep,
+                cross_attn_mask=cross_attn_mask,
                 cross_attention_kwargs=cross_attention_kwargs,
             )
 
@@ -466,6 +468,7 @@ def forward(
         encoder_hidden_states=None,
         timestep=None,
         attention_mask=None,
+        cross_attn_mask: Optional[torch.Tensor] = None,
         cross_attention_kwargs=None,
     ):
         # 1. Self-Attention
@@ -477,6 +480,7 @@ def forward(
             norm_hidden_states,
             encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
             attention_mask=attention_mask,
+            cross_attn_mask=cross_attn_mask if self.only_cross_attention else None,
             **cross_attention_kwargs,
         )
         hidden_states = attn_output + hidden_states
@@ -490,6 +494,7 @@ def forward(
                 norm_hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
                 attention_mask=attention_mask,
+                cross_attn_mask=cross_attn_mask,
                 **cross_attention_kwargs,
             )
             hidden_states = attn_output + hidden_states

src/diffusers/models/cross_attention.py

@@ -173,7 +173,7 @@ def set_subquadratic_attention(
     def set_processor(self, processor: "AttnProcessor"):
         self.processor = processor
 
-    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, **cross_attention_kwargs):
+    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, cross_attn_mask: Optional[torch.Tensor] = None, **cross_attention_kwargs):
         # The `CrossAttention` class can call different attention processors / attention functions
         # here we simply pass along all tensors to the selected processor class
         # For standard processors that are defined here, `**cross_attention_kwargs` is empty
@@ -182,6 +182,7 @@ def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None
             hidden_states,
             encoder_hidden_states=encoder_hidden_states,
             attention_mask=attention_mask,
+            cross_attn_mask=cross_attn_mask,
             **cross_attention_kwargs,
         )
 
@@ -199,11 +200,20 @@ def head_to_batch_dim(self, tensor):
         tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)
         return tensor
 
-    def get_attention_scores(self, query, key, attention_mask=None):
+    def get_attention_scores(self, query, key, attention_mask=None, cross_attn_mask: Optional[torch.Tensor] = None):
         dtype = query.dtype
         if self.upcast_attention:
             query = query.float()
             key = key.float()
+        
+        # haven't defined what to do if both are present
+        if attention_mask is not None:
+            assert cross_attn_mask is None
+        if cross_attn_mask is not None:
+            assert attention_mask is None
+            device = cross_attn_mask.device
+            cross_attn_mask = cross_attn_mask.to('cpu' if device.type == 'mps' else device).repeat_interleave(self.heads, dim=0).to(device).unsqueeze(1)
+            attention_mask = cross_attn_mask
 
         beta = 0 if attention_mask is None else 1
         add = torch.empty(
@@ -242,7 +252,7 @@ def prepare_attention_mask(self, attention_mask, target_length):
 
 
 class CrossAttnProcessor:
-    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, cross_attn_mask: Optional[torch.Tensor] = None):
         batch_size, sequence_length, _ = hidden_states.shape
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
 
@@ -255,7 +265,7 @@ def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=No
         key = attn.head_to_batch_dim(key)
         value = attn.head_to_batch_dim(value)
 
-        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        attention_probs = attn.get_attention_scores(query, key, attention_mask, cross_attn_mask=cross_attn_mask)
         hidden_states = torch.bmm(attention_probs, value)
         hidden_states = attn.batch_to_head_dim(hidden_states)
 

src/diffusers/models/unet_2d_blocks.py

@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 from torch import nn
+from typing import Optional
 
 from .attention import AttentionBlock, DualTransformer2DModel, Transformer2DModel
 from .cross_attention import CrossAttention, CrossAttnAddedKVProcessor
@@ -483,13 +484,14 @@ def __init__(
         self.resnets = nn.ModuleList(resnets)
 
     def forward(
-        self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None
+        self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attn_mask: Optional[torch.Tensor] = None, cross_attention_kwargs=None
     ):
         hidden_states = self.resnets[0](hidden_states, temb)
         for attn, resnet in zip(self.attentions, self.resnets[1:]):
             hidden_states = attn(
                 hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
+                cross_attn_mask=cross_attn_mask,
                 cross_attention_kwargs=cross_attention_kwargs,
             ).sample
             hidden_states = resnet(hidden_states, temb)
@@ -758,7 +760,7 @@ def __init__(
         self.gradient_checkpointing = False
 
     def forward(
-        self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None
+        self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attn_mask: Optional[torch.Tensor] = None, cross_attention_kwargs=None
     ):
         # TODO(Patrick, William) - attention mask is not used
         output_states = ()
@@ -787,6 +789,7 @@ def custom_forward(*inputs):
                 hidden_states = attn(
                     hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
+                    cross_attn_mask=cross_attn_mask,
                     cross_attention_kwargs=cross_attention_kwargs,
                 ).sample
 
@@ -1549,6 +1552,7 @@ def forward(
         cross_attention_kwargs=None,
         upsample_size=None,
         attention_mask=None,
+        cross_attn_mask: Optional[torch.Tensor] = None,
     ):
         # TODO(Patrick, William) - attention mask is not used
         for resnet, attn in zip(self.resnets, self.attentions):
@@ -1580,6 +1584,7 @@ def custom_forward(*inputs):
                 hidden_states = attn(
                     hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
+                    cross_attn_mask=cross_attn_mask,
                     cross_attention_kwargs=cross_attention_kwargs,
                 ).sample
 

src/diffusers/models/unet_2d_condition.py

@@ -354,6 +354,7 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         class_labels: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
+        cross_attn_mask: Optional[torch.Tensor] = None,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
     ) -> Union[UNet2DConditionOutput, Tuple]:
@@ -384,6 +385,9 @@ def forward(
             logger.info("Forward upsample size to force interpolation output size.")
             forward_upsample_size = True
 
+        if cross_attn_mask is not None:
+            cross_attn_mask = (1 - cross_attn_mask.to(sample.dtype)) * -torch.finfo(sample.dtype).max
+
         # prepare attention_mask
         if attention_mask is not None:
             attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
@@ -440,6 +444,7 @@ def forward(
                     temb=emb,
                     encoder_hidden_states=encoder_hidden_states,
                     attention_mask=attention_mask,
+                    cross_attn_mask=cross_attn_mask,
                     cross_attention_kwargs=cross_attention_kwargs,
                 )
             else:
@@ -453,6 +458,7 @@ def forward(
             emb,
             encoder_hidden_states=encoder_hidden_states,
             attention_mask=attention_mask,
+            cross_attn_mask=cross_attn_mask,
             cross_attention_kwargs=cross_attention_kwargs,
         )
 
@@ -477,6 +483,7 @@ def forward(
                     cross_attention_kwargs=cross_attention_kwargs,
                     upsample_size=upsample_size,
                     attention_mask=attention_mask,
+                    cross_attn_mask=cross_attn_mask,
                 )
             else:
                 sample = upsample_block(