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