| 85 | |
| 86 | |
| 87 | class UpResnetBlock1D(nn.Module): |
| 88 | def __init__( |
| 89 | self, |
| 90 | in_channels: int, |
| 91 | out_channels: Optional[int] = None, |
| 92 | num_layers: int = 1, |
| 93 | temb_channels: int = 32, |
| 94 | groups: int = 32, |
| 95 | groups_out: Optional[int] = None, |
| 96 | non_linearity: Optional[str] = None, |
| 97 | time_embedding_norm: str = "default", |
| 98 | output_scale_factor: float = 1.0, |
| 99 | add_upsample: bool = True, |
| 100 | ): |
| 101 | super().__init__() |
| 102 | self.in_channels = in_channels |
| 103 | out_channels = in_channels if out_channels is None else out_channels |
| 104 | self.out_channels = out_channels |
| 105 | self.time_embedding_norm = time_embedding_norm |
| 106 | self.add_upsample = add_upsample |
| 107 | self.output_scale_factor = output_scale_factor |
| 108 | |
| 109 | if groups_out is None: |
| 110 | groups_out = groups |
| 111 | |
| 112 | # there will always be at least one resnet |
| 113 | resnets = [ResidualTemporalBlock1D(2 * in_channels, out_channels, embed_dim=temb_channels)] |
| 114 | |
| 115 | for _ in range(num_layers): |
| 116 | resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=temb_channels)) |
| 117 | |
| 118 | self.resnets = nn.ModuleList(resnets) |
| 119 | |
| 120 | if non_linearity is None: |
| 121 | self.nonlinearity = None |
| 122 | else: |
| 123 | self.nonlinearity = get_activation(non_linearity) |
| 124 | |
| 125 | self.upsample = None |
| 126 | if add_upsample: |
| 127 | self.upsample = Upsample1D(out_channels, use_conv_transpose=True) |
| 128 | |
| 129 | def forward( |
| 130 | self, |
| 131 | hidden_states: torch.Tensor, |
| 132 | res_hidden_states_tuple: Optional[Tuple[torch.Tensor, ...]] = None, |
| 133 | temb: Optional[torch.Tensor] = None, |
| 134 | ) -> torch.Tensor: |
| 135 | if res_hidden_states_tuple is not None: |
| 136 | res_hidden_states = res_hidden_states_tuple[-1] |
| 137 | hidden_states = torch.cat((hidden_states, res_hidden_states), dim=1) |
| 138 | |
| 139 | hidden_states = self.resnets[0](hidden_states, temb) |
| 140 | for resnet in self.resnets[1:]: |
| 141 | hidden_states = resnet(hidden_states, temb) |
| 142 | |
| 143 | if self.nonlinearity is not None: |
| 144 | hidden_states = self.nonlinearity(hidden_states) |