MuseV/musev/models/attention_processor.py

# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""该模型是自定义的attn_processor，实现特殊功能的 Attn功能。
 相对而言，开源代码经常会重新定义Attention 类，
 
 This module implements special AttnProcessor function with custom attn_processor class.
 While other open source code always modify Attention class.
"""
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
from __future__ import annotations

import time
from typing import Any, Callable, Optional
import logging

from einops import rearrange, repeat
import torch
import torch.nn as nn
import torch.nn.functional as F
import xformers
from diffusers.models.lora import LoRACompatibleLinear

from diffusers.utils.torch_utils import maybe_allow_in_graph
from diffusers.models.attention_processor import (
 Attention as DiffusersAttention,
 AttnProcessor,
 AttnProcessor2_0,
)
from ..data.data_util import (
 batch_concat_two_tensor_with_index,
 batch_index_select,
 align_repeat_tensor_single_dim,
 batch_adain_conditioned_tensor,
)

from . import Model_Register

logger = logging.getLogger(__name__) # pylint: disable=invalid-name


@maybe_allow_in_graph
class IPAttention(DiffusersAttention):
 r"""
 Modified Attention class which has special layer, like ip_apadapter_to_k, ip_apadapter_to_v,
 """

 def __init__(
 self,
 query_dim: int,
 cross_attention_dim: int | None = None,
 heads: int = 8,
 dim_head: int = 64,
 dropout: float = 0,
 bias=False,
 upcast_attention: bool = False,
 upcast_softmax: bool = False,
 cross_attention_norm: str | None = None,
 cross_attention_norm_num_groups: int = 32,
 added_kv_proj_dim: int | None = None,
 norm_num_groups: int | None = None,
 spatial_norm_dim: int | None = None,
 out_bias: bool = True,
 scale_qk: bool = True,
 only_cross_attention: bool = False,
 eps: float = 0.00001,
 rescale_output_factor: float = 1,
 residual_connection: bool = False,
 _from_deprecated_attn_block=False,
 processor: AttnProcessor | None = None,
 cross_attn_temporal_cond: bool = False,
 image_scale: float = 1.0,
 ip_adapter_dim: int = None,
 need_t2i_facein: bool = False,
 facein_dim: int = None,
 need_t2i_ip_adapter_face: bool = False,
 ip_adapter_face_dim: int = None,
 ):
 super().__init__(
 query_dim,
 cross_attention_dim,
 heads,
 dim_head,
 dropout,
 bias,
 upcast_attention,
 upcast_softmax,
 cross_attention_norm,
 cross_attention_norm_num_groups,
 added_kv_proj_dim,
 norm_num_groups,
 spatial_norm_dim,
 out_bias,
 scale_qk,
 only_cross_attention,
 eps,
 rescale_output_factor,
 residual_connection,
 _from_deprecated_attn_block,
 processor,
 )
 self.cross_attn_temporal_cond = cross_attn_temporal_cond
 self.image_scale = image_scale
 # 面向首帧的 ip_adapter
 # ip_apdater
 if cross_attn_temporal_cond:
 self.to_k_ip = LoRACompatibleLinear(ip_adapter_dim, query_dim, bias=False)
 self.to_v_ip = LoRACompatibleLinear(ip_adapter_dim, query_dim, bias=False)
 # facein
 self.need_t2i_facein = need_t2i_facein
 self.facein_dim = facein_dim
 if need_t2i_facein:
 raise NotImplementedError("facein")

 # ip_adapter_face
 self.need_t2i_ip_adapter_face = need_t2i_ip_adapter_face
 self.ip_adapter_face_dim = ip_adapter_face_dim
 if need_t2i_ip_adapter_face:
 self.ip_adapter_face_to_k_ip = LoRACompatibleLinear(
 ip_adapter_face_dim, query_dim, bias=False
 )
 self.ip_adapter_face_to_v_ip = LoRACompatibleLinear(
 ip_adapter_face_dim, query_dim, bias=False
 )

 def set_use_memory_efficient_attention_xformers(
 self,
 use_memory_efficient_attention_xformers: bool,
 attention_op: Callable[..., Any] | None = None,
 ):
 if (
 "XFormers" in self.processor.__class__.__name__
 or "IP" in self.processor.__class__.__name__
 ):
 pass
 else:
 return super().set_use_memory_efficient_attention_xformers(
 use_memory_efficient_attention_xformers, attention_op
 )


@Model_Register.register
class BaseIPAttnProcessor(nn.Module):
 print_idx = 0

 def __init__(self, *args, **kwargs) -> None:
 super().__init__(*args, **kwargs)


@Model_Register.register
class T2IReferencenetIPAdapterXFormersAttnProcessor(BaseIPAttnProcessor):
 r"""
 面向 ref_image的 self_attn的 IPAdapter
 """
 print_idx = 0

 def __init__(
 self,
 attention_op: Optional[Callable] = None,
 ):
 super().__init__()

 self.attention_op = attention_op

 def __call__(
 self,
 attn: IPAttention,
 hidden_states: torch.FloatTensor,
 encoder_hidden_states: Optional[torch.FloatTensor] = None,
 attention_mask: Optional[torch.FloatTensor] = None,
 temb: Optional[torch.FloatTensor] = None,
 scale: float = 1.0,
 num_frames: int = None,
 sample_index: torch.LongTensor = None,
 vision_conditon_frames_sample_index: torch.LongTensor = None,
 refer_emb: torch.Tensor = None,
 vision_clip_emb: torch.Tensor = None,
 ip_adapter_scale: float = 1.0,
 face_emb: torch.Tensor = None,
 facein_scale: float = 1.0,
 ip_adapter_face_emb: torch.Tensor = None,
 ip_adapter_face_scale: float = 1.0,
 do_classifier_free_guidance: bool = False,
 ):
 residual = hidden_states

 if attn.spatial_norm is not None:
 hidden_states = attn.spatial_norm(hidden_states, temb)

 input_ndim = hidden_states.ndim

 if input_ndim == 4:
 batch_size, channel, height, width = hidden_states.shape
 hidden_states = hidden_states.view(
 batch_size, channel, height * width
 ).transpose(1, 2)

 batch_size, key_tokens, _ = (
 hidden_states.shape
 if encoder_hidden_states is None
 else encoder_hidden_states.shape
 )

 attention_mask = attn.prepare_attention_mask(
 attention_mask, key_tokens, batch_size
 )
 if attention_mask is not None:
 # expand our mask's singleton query_tokens dimension:
 # [batch*heads, 1, key_tokens] ->
 # [batch*heads, query_tokens, key_tokens]
 # so that it can be added as a bias onto the attention scores that xformers computes:
 # [batch*heads, query_tokens, key_tokens]
 # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
 _, query_tokens, _ = hidden_states.shape
 attention_mask = attention_mask.expand(-1, query_tokens, -1)

 if attn.group_norm is not None:
 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
 1, 2
 )

 query = attn.to_q(hidden_states, scale=scale)

 if encoder_hidden_states is None:
 encoder_hidden_states = hidden_states
 elif attn.norm_cross:
 encoder_hidden_states = attn.norm_encoder_hidden_states(
 encoder_hidden_states
 )
 encoder_hidden_states = align_repeat_tensor_single_dim(
 encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
 )
 key = attn.to_k(encoder_hidden_states, scale=scale)
 value = attn.to_v(encoder_hidden_states, scale=scale)

 # for facein
 if self.print_idx == 0:
 logger.debug(
 f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(face_emb)={type(face_emb)}, facein_scale={facein_scale}"
 )
 if facein_scale > 0 and face_emb is not None:
 raise NotImplementedError("facein")

 query = attn.head_to_batch_dim(query).contiguous()
 key = attn.head_to_batch_dim(key).contiguous()
 value = attn.head_to_batch_dim(value).contiguous()
 hidden_states = xformers.ops.memory_efficient_attention(
 query,
 key,
 value,
 attn_bias=attention_mask,
 op=self.attention_op,
 scale=attn.scale,
 )

 # ip-adapter start
 if self.print_idx == 0:
 logger.debug(
 f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(vision_clip_emb)={type(vision_clip_emb)}"
 )
 if ip_adapter_scale > 0 and vision_clip_emb is not None:
 if self.print_idx == 0:
 logger.debug(
 f"T2I cross_attn, ipadapter, vision_clip_emb={vision_clip_emb.shape}, hidden_states={hidden_states.shape}, batch_size={batch_size}"
 )
 ip_key = attn.to_k_ip(vision_clip_emb)
 ip_value = attn.to_v_ip(vision_clip_emb)
 ip_key = align_repeat_tensor_single_dim(
 ip_key, target_length=batch_size, dim=0
 )
 ip_value = align_repeat_tensor_single_dim(
 ip_value, target_length=batch_size, dim=0
 )
 ip_key = attn.head_to_batch_dim(ip_key).contiguous()
 ip_value = attn.head_to_batch_dim(ip_value).contiguous()
 if self.print_idx == 0:
 logger.debug(
 f"query={query.shape}, ip_key={ip_key.shape}, ip_value={ip_value.shape}"
 )
 # the output of sdp = (batch, num_heads, seq_len, head_dim)
 hidden_states_from_ip = xformers.ops.memory_efficient_attention(
 query,
 ip_key,
 ip_value,
 attn_bias=attention_mask,
 op=self.attention_op,
 scale=attn.scale,
 )
 hidden_states = hidden_states + ip_adapter_scale * hidden_states_from_ip
 # ip-adapter end

 # ip-adapter face start
 if self.print_idx == 0:
 logger.debug(
 f"T2IReferencenetIPAdapterXFormersAttnProcessor,type(ip_adapter_face_emb)={type(ip_adapter_face_emb)}"
 )
 if ip_adapter_face_scale > 0 and ip_adapter_face_emb is not None:
 if self.print_idx == 0:
 logger.debug(
 f"T2I cross_attn, ipadapter face, ip_adapter_face_emb={vision_clip_emb.shape}, hidden_states={hidden_states.shape}, batch_size={batch_size}"
 )
 ip_key = attn.ip_adapter_face_to_k_ip(ip_adapter_face_emb)
 ip_value = attn.ip_adapter_face_to_v_ip(ip_adapter_face_emb)
 ip_key = align_repeat_tensor_single_dim(
 ip_key, target_length=batch_size, dim=0
 )
 ip_value = align_repeat_tensor_single_dim(
 ip_value, target_length=batch_size, dim=0
 )
 ip_key = attn.head_to_batch_dim(ip_key).contiguous()
 ip_value = attn.head_to_batch_dim(ip_value).contiguous()
 if self.print_idx == 0:
 logger.debug(
 f"query={query.shape}, ip_key={ip_key.shape}, ip_value={ip_value.shape}"
 )
 # the output of sdp = (batch, num_heads, seq_len, head_dim)
 hidden_states_from_ip = xformers.ops.memory_efficient_attention(
 query,
 ip_key,
 ip_value,
 attn_bias=attention_mask,
 op=self.attention_op,
 scale=attn.scale,
 )
 hidden_states = (
 hidden_states + ip_adapter_face_scale * hidden_states_from_ip
 )
 # ip-adapter face end

 hidden_states = hidden_states.to(query.dtype)
 hidden_states = attn.batch_to_head_dim(hidden_states)

 # linear proj
 hidden_states = attn.to_out[0](hidden_states, scale=scale)
 # dropout
 hidden_states = attn.to_out[1](hidden_states)

 if input_ndim == 4:
 hidden_states = hidden_states.transpose(-1, -2).reshape(
 batch_size, channel, height, width
 )

 if attn.residual_connection:
 hidden_states = hidden_states + residual

 hidden_states = hidden_states / attn.rescale_output_factor
 self.print_idx += 1
 return hidden_states


@Model_Register.register
class NonParamT2ISelfReferenceXFormersAttnProcessor(BaseIPAttnProcessor):
 r"""
 面向首帧的 referenceonly attn,适用于 T2I的 self_attn
 referenceonly with vis_cond as key, value, in t2i self_attn.
 """
 print_idx = 0

 def __init__(
 self,
 attention_op: Optional[Callable] = None,
 ):
 super().__init__()

 self.attention_op = attention_op

 def __call__(
 self,
 attn: IPAttention,
 hidden_states: torch.FloatTensor,
 encoder_hidden_states: Optional[torch.FloatTensor] = None,
 attention_mask: Optional[torch.FloatTensor] = None,
 temb: Optional[torch.FloatTensor] = None,
 scale: float = 1.0,
 num_frames: int = None,
 sample_index: torch.LongTensor = None,
 vision_conditon_frames_sample_index: torch.LongTensor = None,
 refer_emb: torch.Tensor = None,
 face_emb: torch.Tensor = None,
 vision_clip_emb: torch.Tensor = None,
 ip_adapter_scale: float = 1.0,
 facein_scale: float = 1.0,
 ip_adapter_face_emb: torch.Tensor = None,
 ip_adapter_face_scale: float = 1.0,
 do_classifier_free_guidance: bool = False,
 ):
 residual = hidden_states

 if attn.spatial_norm is not None:
 hidden_states = attn.spatial_norm(hidden_states, temb)

 input_ndim = hidden_states.ndim

 if input_ndim == 4:
 batch_size, channel, height, width = hidden_states.shape
 hidden_states = hidden_states.view(
 batch_size, channel, height * width
 ).transpose(1, 2)

 batch_size, key_tokens, _ = (
 hidden_states.shape
 if encoder_hidden_states is None
 else encoder_hidden_states.shape
 )

 attention_mask = attn.prepare_attention_mask(
 attention_mask, key_tokens, batch_size
 )
 if attention_mask is not None:
 # expand our mask's singleton query_tokens dimension:
 # [batch*heads, 1, key_tokens] ->
 # [batch*heads, query_tokens, key_tokens]
 # so that it can be added as a bias onto the attention scores that xformers computes:
 # [batch*heads, query_tokens, key_tokens]
 # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
 _, query_tokens, _ = hidden_states.shape
 attention_mask = attention_mask.expand(-1, query_tokens, -1)

 # vision_cond in same unet attn start
 if (
 vision_conditon_frames_sample_index is not None and num_frames > 1
 ) or refer_emb is not None:
 batchsize_timesize = hidden_states.shape[0]
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor 0, hidden_states={hidden_states.shape}, vision_conditon_frames_sample_index={vision_conditon_frames_sample_index}"
 )
 encoder_hidden_states = rearrange(
 hidden_states, "(b t) hw c -> b t hw c", t=num_frames
 )
 # if False:
 if vision_conditon_frames_sample_index is not None and num_frames > 1:
 ip_hidden_states = batch_index_select(
 encoder_hidden_states,
 dim=1,
 index=vision_conditon_frames_sample_index,
 ).contiguous()
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor 1, vis_cond referenceonly, encoder_hidden_states={encoder_hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
 )
 #
 ip_hidden_states = rearrange(
 ip_hidden_states, "b t hw c -> b 1 (t hw) c"
 )
 ip_hidden_states = align_repeat_tensor_single_dim(
 ip_hidden_states,
 dim=1,
 target_length=num_frames,
 )
 # b t hw c -> b t hw + hw c
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor 2, vis_cond referenceonly, encoder_hidden_states={encoder_hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
 )
 encoder_hidden_states = torch.concat(
 [encoder_hidden_states, ip_hidden_states], dim=2
 )
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor 3, hidden_states={hidden_states.shape}, ip_hidden_states={ip_hidden_states.shape}"
 )
 # if False:
 if refer_emb is not None: # and num_frames > 1:
 refer_emb = rearrange(refer_emb, "b c t h w->b 1 (t h w) c")
 refer_emb = align_repeat_tensor_single_dim(
 refer_emb, target_length=num_frames, dim=1
 )
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor4, referencenet, encoder_hidden_states={encoder_hidden_states.shape}, refer_emb={refer_emb.shape}"
 )
 encoder_hidden_states = torch.concat(
 [encoder_hidden_states, refer_emb], dim=2
 )
 if self.print_idx == 0:
 logger.debug(
 f"NonParamT2ISelfReferenceXFormersAttnProcessor5, referencenet, encoder_hidden_states={encoder_hidden_states.shape}, refer_emb={refer_emb.shape}"
 )
 encoder_hidden_states = rearrange(
 encoder_hidden_states, "b t hw c -> (b t) hw c"
 )
 # vision_cond in same unet attn end

 if attn.group_norm is not None:
 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
 1, 2
 )

 query = attn.to_q(hidden_states, scale=scale)

 if encoder_hidden_states is None:
 encoder_hidden_states = hidden_states
 elif attn.norm_cross:
 encoder_hidden_states = attn.norm_encoder_hidden_states(
 encoder_hidden_states
 )
 encoder_hidden_states = align_repeat_tensor_single_dim(
 encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
 )
 key = attn.to_k(encoder_hidden_states, scale=scale)
 value = attn.to_v(encoder_hidden_states, scale=scale)

 query = attn.head_to_batch_dim(query).contiguous()
 key = attn.head_to_batch_dim(key).contiguous()
 value = attn.head_to_batch_dim(value).contiguous()

 hidden_states = xformers.ops.memory_efficient_attention(
 query,
 key,
 value,
 attn_bias=attention_mask,
 op=self.attention_op,
 scale=attn.scale,
 )
 hidden_states = hidden_states.to(query.dtype)
 hidden_states = attn.batch_to_head_dim(hidden_states)

 # linear proj
 hidden_states = attn.to_out[0](hidden_states, scale=scale)
 # dropout
 hidden_states = attn.to_out[1](hidden_states)

 if input_ndim == 4:
 hidden_states = hidden_states.transpose(-1, -2).reshape(
 batch_size, channel, height, width
 )

 if attn.residual_connection:
 hidden_states = hidden_states + residual

 hidden_states = hidden_states / attn.rescale_output_factor
 self.print_idx += 1

 return hidden_states


@Model_Register.register
class NonParamReferenceIPXFormersAttnProcessor(
 NonParamT2ISelfReferenceXFormersAttnProcessor
):
 def __init__(self, attention_op: Callable[..., Any] | None = None):
 super().__init__(attention_op)


@maybe_allow_in_graph
class ReferEmbFuseAttention(IPAttention):
 """使用 attention 融合 refernet 中的 emb 到 unet 对应的 latens 中
 # TODO: 目前只支持 bt hw c 的融合，后续考虑增加对 视频 bhw t c、b thw c的融合
 residual_connection: bool = True, 默认， 从不产生影响开始学习

 use attention to fuse referencenet emb into unet latents
 # TODO: by now, only support bt hw c, later consider to support bhw t c, b thw c
 residual_connection: bool = True, default, start from no effect

 Args:
 IPAttention (_type_): _description_
 """

 print_idx = 0

 def __init__(
 self,
 query_dim: int,
 cross_attention_dim: int | None = None,
 heads: int = 8,
 dim_head: int = 64,
 dropout: float = 0,
 bias=False,
 upcast_attention: bool = False,
 upcast_softmax: bool = False,
 cross_attention_norm: str | None = None,
 cross_attention_norm_num_groups: int = 32,
 added_kv_proj_dim: int | None = None,
 norm_num_groups: int | None = None,
 spatial_norm_dim: int | None = None,
 out_bias: bool = True,
 scale_qk: bool = True,
 only_cross_attention: bool = False,
 eps: float = 0.00001,
 rescale_output_factor: float = 1,
 residual_connection: bool = True,
 _from_deprecated_attn_block=False,
 processor: AttnProcessor | None = None,
 cross_attn_temporal_cond: bool = False,
 image_scale: float = 1,
 ):
 super().__init__(
 query_dim,
 cross_attention_dim,
 heads,
 dim_head,
 dropout,
 bias,
 upcast_attention,
 upcast_softmax,
 cross_attention_norm,
 cross_attention_norm_num_groups,
 added_kv_proj_dim,
 norm_num_groups,
 spatial_norm_dim,
 out_bias,
 scale_qk,
 only_cross_attention,
 eps,
 rescale_output_factor,
 residual_connection,
 _from_deprecated_attn_block,
 processor,
 cross_attn_temporal_cond,
 image_scale,
 )
 self.processor = None
 # 配合residual,使一开始不影响之前结果
 nn.init.zeros_(self.to_out[0].weight)
 nn.init.zeros_(self.to_out[0].bias)

 def forward(
 self,
 hidden_states: torch.FloatTensor,
 encoder_hidden_states: Optional[torch.FloatTensor] = None,
 attention_mask: Optional[torch.FloatTensor] = None,
 temb: Optional[torch.FloatTensor] = None,
 scale: float = 1.0,
 num_frames: int = None,
 ) -> torch.Tensor:
 """fuse referencenet emb b c t2 h2 w2 into unet latents b c t1 h1 w1 with attn
 refer to musev/models/attention_processor.py::NonParamT2ISelfReferenceXFormersAttnProcessor

 Args:
 hidden_states (torch.FloatTensor): unet latents, (b t1) c h1 w1
 encoder_hidden_states (Optional[torch.FloatTensor], optional): referencenet emb b c2 t2 h2 w2. Defaults to None.
 attention_mask (Optional[torch.FloatTensor], optional): _description_. Defaults to None.
 temb (Optional[torch.FloatTensor], optional): _description_. Defaults to None.
 scale (float, optional): _description_. Defaults to 1.0.
 num_frames (int, optional): _description_. Defaults to None.

 Returns:
 torch.Tensor: _description_
 """
 residual = hidden_states
 # start
 hidden_states = rearrange(
 hidden_states, "(b t) c h w -> b c t h w", t=num_frames
 )
 batch_size, channel, t1, height, width = hidden_states.shape
 if self.print_idx == 0:
 logger.debug(
 f"hidden_states={hidden_states.shape},encoder_hidden_states={encoder_hidden_states.shape}"
 )
 # concat with hidden_states b c t1 h1 w1 in hw channel into bt (t2 + 1)hw c
 encoder_hidden_states = rearrange(
 encoder_hidden_states, " b c t2 h w-> b (t2 h w) c"
 )
 encoder_hidden_states = repeat(
 encoder_hidden_states, " b t2hw c -> (b t) t2hw c", t=t1
 )
 hidden_states = rearrange(hidden_states, " b c t h w-> (b t) (h w) c")
 # bt (t2+1)hw d
 encoder_hidden_states = torch.concat(
 [encoder_hidden_states, hidden_states], dim=1
 )
 # encoder_hidden_states = align_repeat_tensor_single_dim(
 # encoder_hidden_states, target_length=hidden_states.shape[0], dim=0
 # )
 # end

 if self.spatial_norm is not None:
 hidden_states = self.spatial_norm(hidden_states, temb)

 _, key_tokens, _ = (
 hidden_states.shape
 if encoder_hidden_states is None
 else encoder_hidden_states.shape
 )

 attention_mask = self.prepare_attention_mask(
 attention_mask, key_tokens, batch_size
 )
 if attention_mask is not None:
 # expand our mask's singleton query_tokens dimension:
 # [batch*heads, 1, key_tokens] ->
 # [batch*heads, query_tokens, key_tokens]
 # so that it can be added as a bias onto the attention scores that xformers computes:
 # [batch*heads, query_tokens, key_tokens]
 # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
 _, query_tokens, _ = hidden_states.shape
 attention_mask = attention_mask.expand(-1, query_tokens, -1)

 if self.group_norm is not None:
 hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(
 1, 2
 )

 query = self.to_q(hidden_states, scale=scale)

 if encoder_hidden_states is None:
 encoder_hidden_states = hidden_states
 elif self.norm_cross:
 encoder_hidden_states = self.norm_encoder_hidden_states(
 encoder_hidden_states
 )

 key = self.to_k(encoder_hidden_states, scale=scale)
 value = self.to_v(encoder_hidden_states, scale=scale)

 query = self.head_to_batch_dim(query).contiguous()
 key = self.head_to_batch_dim(key).contiguous()
 value = self.head_to_batch_dim(value).contiguous()

 # query: b t hw d
 # key/value: bt (t1+1)hw d
 hidden_states = xformers.ops.memory_efficient_attention(
 query,
 key,
 value,
 attn_bias=attention_mask,
 scale=self.scale,
 )
 hidden_states = hidden_states.to(query.dtype)
 hidden_states = self.batch_to_head_dim(hidden_states)

 # linear proj
 hidden_states = self.to_out[0](hidden_states, scale=scale)
 # dropout
 hidden_states = self.to_out[1](hidden_states)

 hidden_states = rearrange(
 hidden_states,
 "bt (h w) c-> bt c h w",
 h=height,
 w=width,
 )
 if self.residual_connection:
 hidden_states = hidden_states + residual

 hidden_states = hidden_states / self.rescale_output_factor
 self.print_idx += 1
 return hidden_states