| 83 | |
| 84 | class AE3DConv(torch.nn.Conv2d): |
| 85 | def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs): |
| 86 | super().__init__(in_channels, out_channels, *args, **kwargs) |
| 87 | if isinstance(video_kernel_size, Iterable): |
| 88 | padding = [int(k // 2) for k in video_kernel_size] |
| 89 | else: |
| 90 | padding = int(video_kernel_size // 2) |
| 91 | |
| 92 | self.time_mix_conv = torch.nn.Conv3d( |
| 93 | in_channels=out_channels, |
| 94 | out_channels=out_channels, |
| 95 | kernel_size=video_kernel_size, |
| 96 | padding=padding, |
| 97 | ) |
| 98 | |
| 99 | def forward(self, input, timesteps, skip_video=False): |
| 100 | x = super().forward(input) |