(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, cross_attention_dim=None, norm_num_groups=32, eps=1e-5, need_proj_out=True)
| 126 | class AttentionBlock(torch.nn.Module): |
| 127 | |
| 128 | def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, cross_attention_dim=None, norm_num_groups=32, eps=1e-5, need_proj_out=True): |
| 129 | super().__init__() |
| 130 | inner_dim = num_attention_heads * attention_head_dim |
| 131 | |
| 132 | self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True) |
| 133 | self.proj_in = torch.nn.Linear(in_channels, inner_dim) |
| 134 | |
| 135 | self.transformer_blocks = torch.nn.ModuleList([ |
| 136 | BasicTransformerBlock( |
| 137 | inner_dim, |
| 138 | num_attention_heads, |
| 139 | attention_head_dim, |
| 140 | cross_attention_dim=cross_attention_dim |
| 141 | ) |
| 142 | for d in range(num_layers) |
| 143 | ]) |
| 144 | self.need_proj_out = need_proj_out |
| 145 | if need_proj_out: |
| 146 | self.proj_out = torch.nn.Linear(inner_dim, in_channels) |
| 147 | |
| 148 | def forward( |
| 149 | self, |
nothing calls this directly
no test coverage detected