add group norm type to attention processor cross attention norm

This lets the cross attention norm use both a group norm block and a
layer norm block.

The group norm operates along the channels dimension
and requires input shape (batch size, channels, *) where as the layer norm with a single
`normalized_shape` dimension only operates over the least significant
dimension i.e. (*, channels).

The channels we want to normalize are the hidden dimension of the encoder hidden states.

By convention, the encoder hidden states are always passed as (batch size, sequence
length, hidden states).

This means the layer norm can operate on the tensor without modification, but the group
norm requires flipping the last two dimensions to operate on (batch size, hidden states, sequence length).

All existing attention processors will have the same logic and we can
consolidate it in a helper function `prepare_encoder_hidden_states`

Author: williamberman

src/diffusers/models/attention_processor.py

@@ -56,7 +56,8 @@ def __init__(
         bias=False,
         upcast_attention: bool = False,
         upcast_softmax: bool = False,
-        cross_attention_norm: bool = False,
+        cross_attention_norm: Optional[str] = None,
+        cross_attention_norm_num_groups: int = 32,
         added_kv_proj_dim: Optional[int] = None,
         norm_num_groups: Optional[int] = None,
         out_bias: bool = True,
@@ -68,7 +69,6 @@ def __init__(
         cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
         self.upcast_attention = upcast_attention
         self.upcast_softmax = upcast_softmax
-        self.cross_attention_norm = cross_attention_norm
 
         self.scale = dim_head**-0.5 if scale_qk else 1.0
 
@@ -85,8 +85,28 @@ def __init__(
         else:
             self.group_norm = None
 
-        if cross_attention_norm:
+        if cross_attention_norm is None:
+            self.norm_cross = None
+        elif cross_attention_norm == "layer_norm":
             self.norm_cross = nn.LayerNorm(cross_attention_dim)
+        elif cross_attention_norm == "group_norm":
+            if self.added_kv_proj_dim is not None:
+                # The given `encoder_hidden_states` are initially of shape
+                # (batch_size, seq_len, added_kv_proj_dim) before being projected
+                # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+                # before the projection, so we need to use `added_kv_proj_dim` as
+                # the number of channels for the group norm.
+                norm_cross_num_channels = added_kv_proj_dim
+            else:
+                norm_cross_num_channels = cross_attention_dim
+
+            self.norm_cross = nn.GroupNorm(
+                num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True
+            )
+        else:
+            raise ValueError(
+                f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+            )
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
         self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
@@ -291,6 +311,29 @@ def prepare_attention_mask(self, attention_mask, target_length, batch_size=None)
             attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
         return attention_mask
 
+    def prepare_encoder_hidden_states(self, hidden_states, encoder_hidden_states=None):
+        if encoder_hidden_states is None:
+            return hidden_states
+
+        if self.norm_cross is None:
+            return encoder_hidden_states
+
+        if isinstance(self.norm_cross, nn.LayerNorm):
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+        elif isinstance(self.norm_cross, nn.GroupNorm):
+            # Group norm norms along the channels dimension and expects
+            # input to be in the shape of (N, C, *). In this case, we want
+            # to norm along the hidden dimension, so we need to move
+            # (batch_size, sequence_length, hidden_size) ->
+            # (batch_size, hidden_size, sequence_length)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+        else:
+            assert False
+
+        return encoder_hidden_states
+
 
 class AttnProcessor:
     def __call__(
@@ -306,10 +349,7 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -375,7 +415,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
         query = attn.head_to_batch_dim(query)
 
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
@@ -402,6 +442,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         batch_size, sequence_length, _ = hidden_states.shape
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
@@ -449,10 +490,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         query = attn.to_q(hidden_states)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -493,10 +531,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         query = attn.to_q(hidden_states)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -545,7 +580,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
         query = attn.head_to_batch_dim(query).contiguous()
 
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
@@ -580,10 +615,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         dim = query.shape[-1]
         query = attn.head_to_batch_dim(query)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -630,6 +662,7 @@ def __call__(self, attn: "Attention", hidden_states, encoder_hidden_states=None,
         batch_size, sequence_length, _ = hidden_states.shape
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 

src/diffusers/models/unet_2d_blocks.py

@@ -44,6 +44,7 @@ def get_down_block(
     resnet_time_scale_shift="default",
     resnet_skip_time_act=False,
     resnet_out_scale_factor=1.0,
+    cross_attention_norm=None,
 ):
     down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
     if down_block_type == "DownBlock2D":
@@ -125,6 +126,7 @@ def get_down_block(
             resnet_time_scale_shift=resnet_time_scale_shift,
             skip_time_act=resnet_skip_time_act,
             output_scale_factor=resnet_out_scale_factor,
+            cross_attention_norm=cross_attention_norm,
         )
     elif down_block_type == "SkipDownBlock2D":
         return SkipDownBlock2D(
@@ -222,6 +224,7 @@ def get_up_block(
     resnet_time_scale_shift="default",
     resnet_skip_time_act=False,
     resnet_out_scale_factor=1.0,
+    cross_attention_norm=None,
 ):
     up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
     if up_block_type == "UpBlock2D":
@@ -291,6 +294,7 @@ def get_up_block(
             resnet_time_scale_shift=resnet_time_scale_shift,
             skip_time_act=resnet_skip_time_act,
             output_scale_factor=resnet_out_scale_factor,
+            cross_attention_norm=cross_attention_norm,
         )
     elif up_block_type == "AttnUpBlock2D":
         return AttnUpBlock2D(
@@ -575,6 +579,7 @@ def __init__(
         output_scale_factor=1.0,
         cross_attention_dim=1280,
         skip_time_act=False,
+        cross_attention_norm=None,
     ):
         super().__init__()
 
@@ -614,6 +619,7 @@ def __init__(
                     norm_num_groups=resnet_groups,
                     bias=True,
                     upcast_softmax=True,
+                    cross_attention_norm=cross_attention_norm,
                     processor=AttnAddedKVProcessor(),
                 )
             )
@@ -1356,6 +1362,7 @@ def __init__(
         output_scale_factor=1.0,
         add_downsample=True,
         skip_time_act=False,
+        cross_attention_norm=None,
     ):
         super().__init__()
 
@@ -1394,6 +1401,7 @@ def __init__(
                     norm_num_groups=resnet_groups,
                     bias=True,
                     upcast_softmax=True,
+                    cross_attention_norm=cross_attention_norm,
                     processor=AttnAddedKVProcessor(),
                 )
             )
@@ -1574,7 +1582,7 @@ def __init__(
                     temb_channels=temb_channels,
                     attention_bias=True,
                     add_self_attention=add_self_attention,
-                    cross_attention_norm=True,
+                    cross_attention_norm="layer_norm",
                     group_size=resnet_group_size,
                 )
             )
@@ -2354,6 +2362,7 @@ def __init__(
         output_scale_factor=1.0,
         add_upsample=True,
         skip_time_act=False,
+        cross_attention_norm=None,
     ):
         super().__init__()
         resnets = []
@@ -2393,6 +2402,7 @@ def __init__(
                     norm_num_groups=resnet_groups,
                     bias=True,
                     upcast_softmax=True,
+                    cross_attention_norm=cross_attention_norm,
                     processor=AttnAddedKVProcessor(),
                 )
             )
@@ -2600,7 +2610,7 @@ def __init__(
                     temb_channels=temb_channels,
                     attention_bias=True,
                     add_self_attention=add_self_attention,
-                    cross_attention_norm=True,
+                    cross_attention_norm="layer_norm",
                     upcast_attention=upcast_attention,
                 )
             )
@@ -2695,7 +2705,7 @@ def __init__(
         upcast_attention: bool = False,
         temb_channels: int = 768,  # for ada_group_norm
         add_self_attention: bool = False,
-        cross_attention_norm: bool = False,
+        cross_attention_norm: Optional[str] = None,
         group_size: int = 32,
     ):
         super().__init__()
@@ -2711,7 +2721,7 @@ def __init__(
                 dropout=dropout,
                 bias=attention_bias,
                 cross_attention_dim=None,
-                cross_attention_norm=False,
+                cross_attention_norm=None,
             )
 
         # 2. Cross-Attn

src/diffusers/models/unet_2d_condition.py

@@ -158,6 +158,7 @@ def __init__(
         conv_out_kernel: int = 3,
         projection_class_embeddings_input_dim: Optional[int] = None,
         class_embeddings_concat: bool = False,
+        cross_attention_norm: Optional[str] = None,
     ):
         super().__init__()
 
@@ -311,6 +312,7 @@ def __init__(
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 resnet_skip_time_act=resnet_skip_time_act,
                 resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
             )
             self.down_blocks.append(down_block)
 
@@ -342,6 +344,7 @@ def __init__(
                 resnet_groups=norm_num_groups,
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 skip_time_act=resnet_skip_time_act,
+                cross_attention_norm=cross_attention_norm,
             )
         elif mid_block_type is None:
             self.mid_block = None
@@ -393,6 +396,7 @@ def __init__(
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 resnet_skip_time_act=resnet_skip_time_act,
                 resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
             )
             self.up_blocks.append(up_block)
             prev_output_channel = output_channel

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py

@@ -241,10 +241,7 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py

@@ -63,10 +63,7 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.cross_attention_norm:
-            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)

src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py

@@ -244,6 +244,7 @@ def __init__(
         conv_out_kernel: int = 3,
         projection_class_embeddings_input_dim: Optional[int] = None,
         class_embeddings_concat: bool = False,
+        cross_attention_norm: Optional[str] = None,
     ):
         super().__init__()
 
@@ -403,6 +404,7 @@ def __init__(
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 resnet_skip_time_act=resnet_skip_time_act,
                 resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
             )
             self.down_blocks.append(down_block)
 
@@ -434,6 +436,7 @@ def __init__(
                 resnet_groups=norm_num_groups,
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 skip_time_act=resnet_skip_time_act,
+                cross_attention_norm=cross_attention_norm,
             )
         elif mid_block_type is None:
             self.mid_block = None
@@ -485,6 +488,7 @@ def __init__(
                 resnet_time_scale_shift=resnet_time_scale_shift,
                 resnet_skip_time_act=resnet_skip_time_act,
                 resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
             )
             self.up_blocks.append(up_block)
             prev_output_channel = output_channel
@@ -1463,6 +1467,7 @@ def __init__(
         output_scale_factor=1.0,
         cross_attention_dim=1280,
         skip_time_act=False,
+        cross_attention_norm=None,
     ):
         super().__init__()
 
@@ -1502,6 +1507,7 @@ def __init__(
                     norm_num_groups=resnet_groups,
                     bias=True,
                     upcast_softmax=True,
+                    cross_attention_norm=cross_attention_norm,
                     processor=AttnAddedKVProcessor(),
                 )
             )