| 559 | |
| 560 | |
| 561 | class UpBlock1D(nn.Module): |
| 562 | def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): |
| 563 | super().__init__() |
| 564 | mid_channels = in_channels if mid_channels is None else mid_channels |
| 565 | |
| 566 | resnets = [ |
| 567 | ResConvBlock(2 * in_channels, mid_channels, mid_channels), |
| 568 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 569 | ResConvBlock(mid_channels, mid_channels, out_channels), |
| 570 | ] |
| 571 | |
| 572 | self.resnets = nn.ModuleList(resnets) |
| 573 | self.up = Upsample1d(kernel="cubic") |
| 574 | |
| 575 | def forward( |
| 576 | self, |
| 577 | hidden_states: torch.Tensor, |
| 578 | res_hidden_states_tuple: Tuple[torch.Tensor, ...], |
| 579 | temb: Optional[torch.Tensor] = None, |
| 580 | ) -> torch.Tensor: |
| 581 | res_hidden_states = res_hidden_states_tuple[-1] |
| 582 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 583 | |
| 584 | for resnet in self.resnets: |
| 585 | hidden_states = resnet(hidden_states) |
| 586 | |
| 587 | hidden_states = self.up(hidden_states) |
| 588 | |
| 589 | return hidden_states |
| 590 | |
| 591 | |
| 592 | class UpBlock1DNoSkip(nn.Module): |