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