| 1 | import math |
| 2 | from typing import Tuple, Union |
| 3 | |
| 4 | import torch |
| 5 | from einops import rearrange |
| 6 | from torch import nn |
| 7 | |
| 8 | from .convolution import make_conv_nd |
| 9 | from .enums import PaddingModeType |
| 10 | |
| 11 | |
| 12 | class SpaceToDepthDownsample(nn.Module): |
| 13 | def __init__( |
| 14 | self, |
| 15 | dims: Union[int, Tuple[int, int]], |
| 16 | in_channels: int, |
| 17 | out_channels: int, |
| 18 | stride: Tuple[int, int, int], |
| 19 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 20 | ): |
| 21 | super().__init__() |
| 22 | self.stride = stride |
| 23 | self.group_size = in_channels * math.prod(stride) // out_channels |
| 24 | self.conv = make_conv_nd( |
| 25 | dims=dims, |
| 26 | in_channels=in_channels, |
| 27 | out_channels=out_channels // math.prod(stride), |
| 28 | kernel_size=3, |
| 29 | stride=1, |
| 30 | causal=True, |
| 31 | spatial_padding_mode=spatial_padding_mode, |
| 32 | ) |
| 33 | |
| 34 | def forward( |
| 35 | self, |
| 36 | x: torch.Tensor, |
| 37 | causal: bool = True, |
| 38 | ) -> torch.Tensor: |
| 39 | if self.stride[0] == 2: |
| 40 | x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding |
| 41 | |
| 42 | # skip connection |
| 43 | x_in = rearrange( |
| 44 | x, |
| 45 | "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", |
| 46 | p1=self.stride[0], |
| 47 | p2=self.stride[1], |
| 48 | p3=self.stride[2], |
| 49 | ) |
| 50 | x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size) |
| 51 | x_in = x_in.mean(dim=2) |
| 52 | |
| 53 | # conv |
| 54 | x = self.conv(x, causal=causal) |
| 55 | x = rearrange( |
| 56 | x, |
| 57 | "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", |
| 58 | p1=self.stride[0], |
| 59 | p2=self.stride[1], |
| 60 | p3=self.stride[2], |
| 61 | ) |
| 62 | |
| 63 | x = x + x_in |
| 64 | |
| 65 | return x |
| 66 | |
| 67 | |
| 68 | class DepthToSpaceUpsample(nn.Module): |
| 69 | def __init__( |
| 70 | self, |
| 71 | dims: int | Tuple[int, int], |
| 72 | in_channels: int, |
| 73 | stride: Tuple[int, int, int], |
| 74 | residual: bool = False, |
| 75 | out_channels_reduction_factor: int = 1, |
| 76 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 77 | ): |
| 78 | super().__init__() |
| 79 | self.stride = stride |
| 80 | self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor |
| 81 | self.conv = make_conv_nd( |
| 82 | dims=dims, |
| 83 | in_channels=in_channels, |
| 84 | out_channels=self.out_channels, |
| 85 | kernel_size=3, |
| 86 | stride=1, |
| 87 | causal=True, |
| 88 | spatial_padding_mode=spatial_padding_mode, |
| 89 | ) |
| 90 | self.residual = residual |
| 91 | self.out_channels_reduction_factor = out_channels_reduction_factor |
| 92 | |
| 93 | def forward( |
| 94 | self, |
| 95 | x: torch.Tensor, |
| 96 | causal: bool = True, |
| 97 | ) -> torch.Tensor: |
| 98 | if self.residual: |
| 99 | # Reshape and duplicate the input to match the output shape |
| 100 | x_in = rearrange( |
| 101 | x, |
| 102 | "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", |
| 103 | p1=self.stride[0], |
| 104 | p2=self.stride[1], |
| 105 | p3=self.stride[2], |
| 106 | ) |
| 107 | num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor |
| 108 | x_in = x_in.repeat(1, num_repeat, 1, 1, 1) |
| 109 | if self.stride[0] == 2: |
| 110 | x_in = x_in[:, :, 1:, :, :] |
| 111 | x = self.conv(x, causal=causal) |
| 112 | x = rearrange( |
| 113 | x, |
| 114 | "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", |
| 115 | p1=self.stride[0], |
| 116 | p2=self.stride[1], |
| 117 | p3=self.stride[2], |
| 118 | ) |
| 119 | if self.stride[0] == 2: |
| 120 | x = x[:, :, 1:, :, :] |
| 121 | if self.residual: |
| 122 | x = x + x_in |
| 123 | return x |
| 124 |