| 223 | |
| 224 | |
| 225 | class OutConv1DBlock(nn.Module): |
| 226 | def __init__(self, num_groups_out: int, out_channels: int, embed_dim: int, act_fn: str): |
| 227 | super().__init__() |
| 228 | self.final_conv1d_1 = nn.Conv1d(embed_dim, embed_dim, 5, padding=2) |
| 229 | self.final_conv1d_gn = nn.GroupNorm(num_groups_out, embed_dim) |
| 230 | self.final_conv1d_act = get_activation(act_fn) |
| 231 | self.final_conv1d_2 = nn.Conv1d(embed_dim, out_channels, 1) |
| 232 | |
| 233 | def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: |
| 234 | hidden_states = self.final_conv1d_1(hidden_states) |
| 235 | hidden_states = rearrange_dims(hidden_states) |
| 236 | hidden_states = self.final_conv1d_gn(hidden_states) |
| 237 | hidden_states = rearrange_dims(hidden_states) |
| 238 | hidden_states = self.final_conv1d_act(hidden_states) |
| 239 | hidden_states = self.final_conv1d_2(hidden_states) |
| 240 | return hidden_states |
| 241 | |
| 242 | |
| 243 | class OutValueFunctionBlock(nn.Module): |