r""" A cross attention layer. Parameters: query_dim (`int`): The number of channels in the query. cross_attention_dim (`int`, *optional*): The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. heads (`int`
| 39 | |
| 40 | @maybe_allow_in_graph |
| 41 | class Attention(nn.Module): |
| 42 | r""" |
| 43 | A cross attention layer. |
| 44 | |
| 45 | Parameters: |
| 46 | query_dim (`int`): |
| 47 | The number of channels in the query. |
| 48 | cross_attention_dim (`int`, *optional*): |
| 49 | The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. |
| 50 | heads (`int`, *optional*, defaults to 8): |
| 51 | The number of heads to use for multi-head attention. |
| 52 | kv_heads (`int`, *optional*, defaults to `None`): |
| 53 | The number of key and value heads to use for multi-head attention. Defaults to `heads`. If |
| 54 | `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi |
| 55 | Query Attention (MQA) otherwise GQA is used. |
| 56 | dim_head (`int`, *optional*, defaults to 64): |
| 57 | The number of channels in each head. |
| 58 | dropout (`float`, *optional*, defaults to 0.0): |
| 59 | The dropout probability to use. |
| 60 | bias (`bool`, *optional*, defaults to False): |
| 61 | Set to `True` for the query, key, and value linear layers to contain a bias parameter. |
| 62 | upcast_attention (`bool`, *optional*, defaults to False): |
| 63 | Set to `True` to upcast the attention computation to `float32`. |
| 64 | upcast_softmax (`bool`, *optional*, defaults to False): |
| 65 | Set to `True` to upcast the softmax computation to `float32`. |
| 66 | cross_attention_norm (`str`, *optional*, defaults to `None`): |
| 67 | The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`. |
| 68 | cross_attention_norm_num_groups (`int`, *optional*, defaults to 32): |
| 69 | The number of groups to use for the group norm in the cross attention. |
| 70 | added_kv_proj_dim (`int`, *optional*, defaults to `None`): |
| 71 | The number of channels to use for the added key and value projections. If `None`, no projection is used. |
| 72 | norm_num_groups (`int`, *optional*, defaults to `None`): |
| 73 | The number of groups to use for the group norm in the attention. |
| 74 | spatial_norm_dim (`int`, *optional*, defaults to `None`): |
| 75 | The number of channels to use for the spatial normalization. |
| 76 | out_bias (`bool`, *optional*, defaults to `True`): |
| 77 | Set to `True` to use a bias in the output linear layer. |
| 78 | scale_qk (`bool`, *optional*, defaults to `True`): |
| 79 | Set to `True` to scale the query and key by `1 / sqrt(dim_head)`. |
| 80 | only_cross_attention (`bool`, *optional*, defaults to `False`): |
| 81 | Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if |
| 82 | `added_kv_proj_dim` is not `None`. |
| 83 | eps (`float`, *optional*, defaults to 1e-5): |
| 84 | An additional value added to the denominator in group normalization that is used for numerical stability. |
| 85 | rescale_output_factor (`float`, *optional*, defaults to 1.0): |
| 86 | A factor to rescale the output by dividing it with this value. |
| 87 | residual_connection (`bool`, *optional*, defaults to `False`): |
| 88 | Set to `True` to add the residual connection to the output. |
| 89 | _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`): |
| 90 | Set to `True` if the attention block is loaded from a deprecated state dict. |
| 91 | processor (`AttnProcessor`, *optional*, defaults to `None`): |
| 92 | The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and |
| 93 | `AttnProcessor` otherwise. |
| 94 | """ |
| 95 | |
| 96 | def __init__( |
| 97 | self, |
| 98 | query_dim: int, |
no outgoing calls