| 71 | return x.reshape(-1, C, H, W) |
| 72 | |
| 73 | class PixelShuffle3d(nn.Module): |
| 74 | def __init__(self, ff, hh, ww): |
| 75 | super().__init__() |
| 76 | self.ff = ff |
| 77 | self.hh = hh |
| 78 | self.ww = ww |
| 79 | def forward(self, x): |
| 80 | # x: (B, C, F, H, W) |
| 81 | B, C, F, H, W = x.shape |
| 82 | if F % self.ff != 0: |
| 83 | first_frame = x[:, :, 0:1, :, :].repeat(1, 1, self.ff - F % self.ff, 1, 1) |
| 84 | x = torch.cat([first_frame, x], dim=2) |
| 85 | return rearrange( |
| 86 | x, |
| 87 | 'b c (f ff) (h hh) (w ww) -> b (c ff hh ww) f h w', |
| 88 | ff=self.ff, hh=self.hh, ww=self.ww |
| 89 | ).transpose(1, 2) |
| 90 | |
| 91 | # ---------------------------- |
| 92 | # Generic NTCHW graph executor (kept; used by decoder) |