| 5 | |
| 6 | |
| 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 |
| 29 | residual = hidden_states |
| 30 | |
| 31 | hidden_states = self.norm(hidden_states) |
| 32 | inner_dim = hidden_states.shape[1] |
| 33 | hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) |
| 34 | |
| 35 | for block in self.transformer_blocks: |
| 36 | hidden_states = block(hidden_states) |
| 37 | |
| 38 | hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() |
| 39 | hidden_states = hidden_states + residual |
| 40 | |
| 41 | return hidden_states, time_emb, text_emb, res_stack |
| 42 | |
| 43 | |
| 44 | class SDVAEDecoder(torch.nn.Module): |