(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
cross_attention_dim: Optional[int] = None,
activation_fn: str = "geglu",
num_embeds_ada_norm: Optional[int] = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
upcast_attention: bool = False,
norm_elementwise_affine: bool = True,
norm_type: str = "layer_norm",
final_dropout: bool = False,
)
| 45 | """ |
| 46 | |
| 47 | def __init__( |
| 48 | self, |
| 49 | dim: int, |
| 50 | num_attention_heads: int, |
| 51 | attention_head_dim: int, |
| 52 | dropout=0.0, |
| 53 | cross_attention_dim: Optional[int] = None, |
| 54 | activation_fn: str = "geglu", |
| 55 | num_embeds_ada_norm: Optional[int] = None, |
| 56 | attention_bias: bool = False, |
| 57 | only_cross_attention: bool = False, |
| 58 | double_self_attention: bool = False, |
| 59 | upcast_attention: bool = False, |
| 60 | norm_elementwise_affine: bool = True, |
| 61 | norm_type: str = "layer_norm", |
| 62 | final_dropout: bool = False, |
| 63 | ): |
| 64 | super().__init__() |
| 65 | self.only_cross_attention = only_cross_attention |
| 66 | |
| 67 | self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" |
| 68 | self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" |
| 69 | |
| 70 | if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: |
| 71 | raise ValueError( |
| 72 | f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" |
| 73 | f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." |
| 74 | ) |
| 75 | |
| 76 | # Define 3 blocks. Each block has its own normalization layer. |
| 77 | # 1. Self-Attn |
| 78 | if self.use_ada_layer_norm: |
| 79 | self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) |
| 80 | elif self.use_ada_layer_norm_zero: |
| 81 | self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) |
| 82 | else: |
| 83 | self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) |
| 84 | self.attn1 = Attention( |
| 85 | query_dim=dim, |
| 86 | heads=num_attention_heads, |
| 87 | dim_head=attention_head_dim, |
| 88 | dropout=dropout, |
| 89 | bias=attention_bias, |
| 90 | cross_attention_dim=cross_attention_dim if only_cross_attention else None, |
| 91 | upcast_attention=upcast_attention, |
| 92 | ) |
| 93 | |
| 94 | # 2. Cross-Attn |
| 95 | if cross_attention_dim is not None or double_self_attention: |
| 96 | # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. |
| 97 | # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during |
| 98 | # the second cross attention block. |
| 99 | self.norm2 = ( |
| 100 | AdaLayerNorm(dim, num_embeds_ada_norm) |
| 101 | if self.use_ada_layer_norm |
| 102 | else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) |
| 103 | ) |
| 104 | self.attn2 = Attention( |
nothing calls this directly
no test coverage detected