| 7 | class VAEAttentionBlock(torch.nn.Module): |
| 8 | |
| 9 | def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, norm_num_groups=32, eps=1e-5): |
| 10 | super().__init__() |
| 11 | inner_dim = num_attention_heads * attention_head_dim |
| 12 | |
| 13 | self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True) |
| 14 | |
| 15 | self.transformer_blocks = torch.nn.ModuleList([ |
| 16 | Attention( |
| 17 | inner_dim, |
| 18 | num_attention_heads, |
| 19 | attention_head_dim, |
| 20 | bias_q=True, |
| 21 | bias_kv=True, |
| 22 | bias_out=True |
| 23 | ) |
| 24 | for d in range(num_layers) |
| 25 | ]) |
| 26 | |
| 27 | def forward(self, hidden_states, time_emb, text_emb, res_stack): |
| 28 | batch, _, height, width = hidden_states.shape |