| 521 | |
| 522 | |
| 523 | class AttnUpBlock1D(nn.Module): |
| 524 | def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): |
| 525 | super().__init__() |
| 526 | mid_channels = out_channels if mid_channels is None else mid_channels |
| 527 | |
| 528 | resnets = [ |
| 529 | ResConvBlock(2 * in_channels, mid_channels, mid_channels), |
| 530 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 531 | ResConvBlock(mid_channels, mid_channels, out_channels), |
| 532 | ] |
| 533 | attentions = [ |
| 534 | SelfAttention1d(mid_channels, mid_channels // 32), |
| 535 | SelfAttention1d(mid_channels, mid_channels // 32), |
| 536 | SelfAttention1d(out_channels, out_channels // 32), |
| 537 | ] |
| 538 | |
| 539 | self.attentions = nn.ModuleList(attentions) |
| 540 | self.resnets = nn.ModuleList(resnets) |
| 541 | self.up = Upsample1d(kernel="cubic") |
| 542 | |
| 543 | def forward( |
| 544 | self, |
| 545 | hidden_states: torch.Tensor, |
| 546 | res_hidden_states_tuple: Tuple[torch.Tensor, ...], |
| 547 | temb: Optional[torch.Tensor] = None, |
| 548 | ) -> torch.Tensor: |
| 549 | res_hidden_states = res_hidden_states_tuple[-1] |
| 550 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 551 | |
| 552 | for resnet, attn in zip(self.resnets, self.attentions): |
| 553 | hidden_states = resnet(hidden_states) |
| 554 | hidden_states = attn(hidden_states) |
| 555 | |
| 556 | hidden_states = self.up(hidden_states) |
| 557 | |
| 558 | return hidden_states |
| 559 | |
| 560 | |
| 561 | class UpBlock1D(nn.Module): |