| 447 | |
| 448 | |
| 449 | class AttnDownBlock1D(nn.Module): |
| 450 | def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): |
| 451 | super().__init__() |
| 452 | mid_channels = out_channels if mid_channels is None else mid_channels |
| 453 | |
| 454 | self.down = Downsample1d("cubic") |
| 455 | resnets = [ |
| 456 | ResConvBlock(in_channels, mid_channels, mid_channels), |
| 457 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 458 | ResConvBlock(mid_channels, mid_channels, out_channels), |
| 459 | ] |
| 460 | attentions = [ |
| 461 | SelfAttention1d(mid_channels, mid_channels // 32), |
| 462 | SelfAttention1d(mid_channels, mid_channels // 32), |
| 463 | SelfAttention1d(out_channels, out_channels // 32), |
| 464 | ] |
| 465 | |
| 466 | self.attentions = nn.ModuleList(attentions) |
| 467 | self.resnets = nn.ModuleList(resnets) |
| 468 | |
| 469 | def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: |
| 470 | hidden_states = self.down(hidden_states) |
| 471 | |
| 472 | for resnet, attn in zip(self.resnets, self.attentions): |
| 473 | hidden_states = resnet(hidden_states) |
| 474 | hidden_states = attn(hidden_states) |
| 475 | |
| 476 | return hidden_states, (hidden_states,) |
| 477 | |
| 478 | |
| 479 | class DownBlock1D(nn.Module): |