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