| 436 | return x |
| 437 | |
| 438 | class CausalConvAfterNorm(CausalConv): |
| 439 | def __init__(self, |
| 440 | chan_in, |
| 441 | chan_out, |
| 442 | kernel_size, |
| 443 | **kwargs |
| 444 | ): |
| 445 | super().__init__( |
| 446 | chan_in, chan_out, kernel_size, **kwargs) |
| 447 | |
| 448 | if self.time_causal_padding == (1, 1, 1, 1, 2, 0): |
| 449 | self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, padding=(0, 1, 1), **kwargs) |
| 450 | else: |
| 451 | self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) |
| 452 | self.is_first_run = True |
| 453 | |
| 454 | def forward(self, x, is_init=True, residual=None): |
| 455 | if self.is_first_run: |
| 456 | self.is_first_run = False |
| 457 | |
| 458 | if self.time_causal_padding == (1, 1, 1, 1, 2, 0): |
| 459 | pass |
| 460 | else: |
| 461 | x = nn.functional.pad(x, self.time_causal_padding).contiguous() |
| 462 | |
| 463 | x = base_conv3d_channel_last(x, self.conv, residual=residual) |
| 464 | return x |
| 465 | |
| 466 | class AttnBlock(nn.Module): |
| 467 | def __init__(self, |