| 477 | |
| 478 | |
| 479 | class DownBlock1D(nn.Module): |
| 480 | def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): |
| 481 | super().__init__() |
| 482 | mid_channels = out_channels if mid_channels is None else mid_channels |
| 483 | |
| 484 | self.down = Downsample1d("cubic") |
| 485 | resnets = [ |
| 486 | ResConvBlock(in_channels, mid_channels, mid_channels), |
| 487 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 488 | ResConvBlock(mid_channels, mid_channels, out_channels), |
| 489 | ] |
| 490 | |
| 491 | self.resnets = nn.ModuleList(resnets) |
| 492 | |
| 493 | def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: |
| 494 | hidden_states = self.down(hidden_states) |
| 495 | |
| 496 | for resnet in self.resnets: |
| 497 | hidden_states = resnet(hidden_states) |
| 498 | |
| 499 | return hidden_states, (hidden_states,) |
| 500 | |
| 501 | |
| 502 | class DownBlock1DNoSkip(nn.Module): |