| 22 | |
| 23 | |
| 24 | class DownResnetBlock1D(nn.Module): |
| 25 | def __init__( |
| 26 | self, |
| 27 | in_channels: int, |
| 28 | out_channels: int | None = None, |
| 29 | num_layers: int = 1, |
| 30 | conv_shortcut: bool = False, |
| 31 | temb_channels: int = 32, |
| 32 | groups: int = 32, |
| 33 | groups_out: int | None = None, |
| 34 | non_linearity: str | None = None, |
| 35 | time_embedding_norm: str = "default", |
| 36 | output_scale_factor: float = 1.0, |
| 37 | add_downsample: bool = True, |
| 38 | ): |
| 39 | super().__init__() |
| 40 | self.in_channels = in_channels |
| 41 | out_channels = in_channels if out_channels is None else out_channels |
| 42 | self.out_channels = out_channels |
| 43 | self.use_conv_shortcut = conv_shortcut |
| 44 | self.time_embedding_norm = time_embedding_norm |
| 45 | self.add_downsample = add_downsample |
| 46 | self.output_scale_factor = output_scale_factor |
| 47 | |
| 48 | if groups_out is None: |
| 49 | groups_out = groups |
| 50 | |
| 51 | # there will always be at least one resnet |
| 52 | resnets = [ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=temb_channels)] |
| 53 | |
| 54 | for _ in range(num_layers): |
| 55 | resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=temb_channels)) |
| 56 | |
| 57 | self.resnets = nn.ModuleList(resnets) |
| 58 | |
| 59 | if non_linearity is None: |
| 60 | self.nonlinearity = None |
| 61 | else: |
| 62 | self.nonlinearity = get_activation(non_linearity) |
| 63 | |
| 64 | self.downsample = None |
| 65 | if add_downsample: |
| 66 | self.downsample = Downsample1D(out_channels, use_conv=True, padding=1) |
| 67 | |
| 68 | def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: |
| 69 | output_states = () |
| 70 | |
| 71 | hidden_states = self.resnets[0](hidden_states, temb) |
| 72 | for resnet in self.resnets[1:]: |
| 73 | hidden_states = resnet(hidden_states, temb) |
| 74 | |
| 75 | output_states += (hidden_states,) |
| 76 | |
| 77 | if self.nonlinearity is not None: |
| 78 | hidden_states = self.nonlinearity(hidden_states) |
| 79 | |
| 80 | if self.downsample is not None: |
| 81 | hidden_states = self.downsample(hidden_states) |
no outgoing calls
no test coverage detected
searching dependent graphs…