| 500 | |
| 501 | |
| 502 | class DownBlock1DNoSkip(nn.Module): |
| 503 | def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): |
| 504 | super().__init__() |
| 505 | mid_channels = out_channels if mid_channels is None else mid_channels |
| 506 | |
| 507 | resnets = [ |
| 508 | ResConvBlock(in_channels, mid_channels, mid_channels), |
| 509 | ResConvBlock(mid_channels, mid_channels, mid_channels), |
| 510 | ResConvBlock(mid_channels, mid_channels, out_channels), |
| 511 | ] |
| 512 | |
| 513 | self.resnets = nn.ModuleList(resnets) |
| 514 | |
| 515 | def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: |
| 516 | hidden_states = torch.cat([hidden_states, temb], dim=1) |
| 517 | for resnet in self.resnets: |
| 518 | hidden_states = resnet(hidden_states) |
| 519 | |
| 520 | return hidden_states, (hidden_states,) |
| 521 | |
| 522 | |
| 523 | class AttnUpBlock1D(nn.Module): |