| 1 | import torch |
| 2 | from einops import rearrange |
| 3 | |
| 4 | |
| 5 | class PixelShuffleND(torch.nn.Module): |
| 6 | """ |
| 7 | N-dimensional pixel shuffle operation for upsampling tensors. |
| 8 | Args: |
| 9 | dims (int): Number of dimensions to apply pixel shuffle to. |
| 10 | - 1: Temporal (e.g., frames) |
| 11 | - 2: Spatial (e.g., height and width) |
| 12 | - 3: Spatiotemporal (e.g., depth, height, width) |
| 13 | upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension. |
| 14 | For dims=1, only the first value is used. |
| 15 | For dims=2, the first two values are used. |
| 16 | For dims=3, all three values are used. |
| 17 | The input tensor is rearranged so that the channel dimension is split into |
| 18 | smaller channels and upscaling factors, and the upscaling factors are moved |
| 19 | into the corresponding spatial/temporal dimensions. |
| 20 | Note: |
| 21 | This operation is equivalent to the patchifier operation in for the models. Consider |
| 22 | using this class instead. |
| 23 | """ |
| 24 | |
| 25 | def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)): |
| 26 | super().__init__() |
| 27 | assert dims in [1, 2, 3], "dims must be 1, 2, or 3" |
| 28 | self.dims = dims |
| 29 | self.upscale_factors = upscale_factors |
| 30 | |
| 31 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 32 | if self.dims == 3: |
| 33 | return rearrange( |
| 34 | x, |
| 35 | "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", |
| 36 | p1=self.upscale_factors[0], |
| 37 | p2=self.upscale_factors[1], |
| 38 | p3=self.upscale_factors[2], |
| 39 | ) |
| 40 | elif self.dims == 2: |
| 41 | return rearrange( |
| 42 | x, |
| 43 | "b (c p1 p2) h w -> b c (h p1) (w p2)", |
| 44 | p1=self.upscale_factors[0], |
| 45 | p2=self.upscale_factors[1], |
| 46 | ) |
| 47 | elif self.dims == 1: |
| 48 | return rearrange( |
| 49 | x, |
| 50 | "b (c p1) f h w -> b c (f p1) h w", |
| 51 | p1=self.upscale_factors[0], |
| 52 | ) |
| 53 | else: |
| 54 | raise ValueError(f"Unsupported dims: {self.dims}") |
| 55 |