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`, *o
| 34 | |
| 35 | @maybe_allow_in_graph |
| 36 | class Attention(nn.Module): |
| 37 | r""" |
| 38 | A cross attention layer. |
| 39 | |
| 40 | Parameters: |
| 41 | query_dim (`int`): The number of channels in the query. |
| 42 | cross_attention_dim (`int`, *optional*): |
| 43 | The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. |
| 44 | heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention. |
| 45 | dim_head (`int`, *optional*, defaults to 64): The number of channels in each head. |
| 46 | dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. |
| 47 | bias (`bool`, *optional*, defaults to False): |
| 48 | Set to `True` for the query, key, and value linear layers to contain a bias parameter. |
| 49 | """ |
| 50 | |
| 51 | def __init__( |
| 52 | self, |
| 53 | query_dim: int, |
| 54 | cross_attention_dim: Optional[int] = None, |
| 55 | heads: int = 8, |
| 56 | dim_head: int = 64, |
| 57 | dropout: float = 0.0, |
| 58 | bias=False, |
| 59 | upcast_attention: bool = False, |
| 60 | upcast_softmax: bool = False, |
| 61 | cross_attention_norm: Optional[str] = None, |
| 62 | cross_attention_norm_num_groups: int = 32, |
| 63 | added_kv_proj_dim: Optional[int] = None, |
| 64 | norm_num_groups: Optional[int] = None, |
| 65 | spatial_norm_dim: Optional[int] = None, |
| 66 | out_bias: bool = True, |
| 67 | scale_qk: bool = True, |
| 68 | only_cross_attention: bool = False, |
| 69 | eps: float = 1e-5, |
| 70 | rescale_output_factor: float = 1.0, |
| 71 | residual_connection: bool = False, |
| 72 | _from_deprecated_attn_block=False, |
| 73 | processor: Optional["AttnProcessor"] = None, |
| 74 | ): |
| 75 | super().__init__() |
| 76 | inner_dim = dim_head * heads |
| 77 | cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim |
| 78 | self.upcast_attention = upcast_attention |
| 79 | self.upcast_softmax = upcast_softmax |
| 80 | self.rescale_output_factor = rescale_output_factor |
| 81 | self.residual_connection = residual_connection |
| 82 | |
| 83 | # we make use of this private variable to know whether this class is loaded |
| 84 | # with an deprecated state dict so that we can convert it on the fly |
| 85 | self._from_deprecated_attn_block = _from_deprecated_attn_block |
| 86 | |
| 87 | self.scale_qk = scale_qk |
| 88 | self.scale = dim_head**-0.5 if self.scale_qk else 1.0 |
| 89 | |
| 90 | self.heads = heads |
| 91 | # for slice_size > 0 the attention score computation |
| 92 | # is split across the batch axis to save memory |
| 93 | # You can set slice_size with `set_attention_slice` |
no outgoing calls
no test coverage detected