| from torch import nn |
| from einops import rearrange |
| import torch as t |
|
|
| class Patch(nn.Module): |
| def __init__(self, in_channels=3, out_channels=64, patch_size=2): |
| super().__init__() |
| self.patch_size = patch_size |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
| dim = out_channels |
| if dim % 32 == 0 and dim > 32: |
| self.init_conv_seq = nn.Sequential( |
| nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), |
| nn.SiLU(), |
| nn.GroupNorm(32, dim // 2), |
| nn.Conv2d(dim // 2, dim // 2, kernel_size=5, padding=2, stride=1), |
| nn.SiLU(), |
| nn.GroupNorm(32, dim // 2), |
| ) |
| else: |
| self.init_conv_seq = nn.Sequential( |
| nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), |
| nn.SiLU(), |
| nn.Conv2d(dim // 2, dim // 2, kernel_size=5, padding=2, stride=1), |
| nn.SiLU(), |
| ) |
|
|
| self.x_embedder = nn.Linear(patch_size * patch_size * dim // 2, dim, bias=True) |
| nn.init.constant_(self.x_embedder.bias, 0) |
|
|
| def forward(self, x): |
| batch, dur, c, h, w = x.shape |
| x = x.reshape(-1, c, h, w) |
| x = self.init_conv_seq(x) |
| x = self.patchify(x) |
| x = self.x_embedder(x) |
| x = x.reshape(batch, dur, -1, self.out_channels) |
| return x |
|
|
| def patchify(self, x): |
| B, C, H, W = x.size() |
| x = x.view( |
| B, |
| C, |
| H // self.patch_size, |
| self.patch_size, |
| W // self.patch_size, |
| self.patch_size, |
| ) |
| x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2) |
| return x |
|
|
| class UnPatch(nn.Module): |
| def __init__(self, height, width, in_channels=64, out_channels=3, patch_size=2): |
| super().__init__() |
| self.width = width |
| self.height = height |
| self.patch_size = patch_size |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
| self.unpatch = nn.Linear(in_channels, out_channels*patch_size**2) |
|
|
| def forward(self, x): |
| x = self.unpatch(x) |
| batch, dur, seq, d = x.shape |
| x = x.reshape(-1, seq, d) |
| x = self.unpatchify(x) |
| x = x.reshape(batch, dur, self.out_channels, self.height, self.width) |
| return x |
| |
| def unpatchify(self, x): |
| c = self.out_channels |
| p = self.patch_size |
| h = self.height // p |
| w = self.width // p |
| x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) |
| x = t.einsum("nhwpqc->nchpwq", x) |
| imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p)) |
| return imgs |
