| 590 | |
| 591 | |
| 592 | class UpBlock1DNoSkip(nn.Module): |
| 593 | def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): |
| 594 | super().__init__() |
| 595 | mid_channels = in_channels if mid_channels is None else mid_channels |
| 596 | |
| 597 | resnets = [ |
| 598 | ResConvBlock(2 * in_channels, mid_channels, mid_channels), |
| 599 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 600 | ResConvBlock(mid_channels, mid_channels, out_channels, is_last=True), |
| 601 | ] |
| 602 | |
| 603 | self.resnets = nn.ModuleList(resnets) |
| 604 | |
| 605 | def forward( |
| 606 | self, |
| 607 | hidden_states: torch.Tensor, |
| 608 | res_hidden_states_tuple: Tuple[torch.Tensor, ...], |
| 609 | temb: Optional[torch.Tensor] = None, |
| 610 | ) -> torch.Tensor: |
| 611 | res_hidden_states = res_hidden_states_tuple[-1] |
| 612 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 613 | |
| 614 | for resnet in self.resnets: |
| 615 | hidden_states = resnet(hidden_states) |
| 616 | |
| 617 | return hidden_states |
| 618 | |
| 619 | |
| 620 | DownBlockType = Union[DownResnetBlock1D, DownBlock1D, AttnDownBlock1D, DownBlock1DNoSkip] |