Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| from einops import rearrange | |
| from models.positional_encoding import build_spatial_only_pe | |
| class PatchEmbedding(nn.Module): | |
| def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128): | |
| super().__init__() | |
| H, W = frame_size | |
| self.frame_size = frame_size | |
| self.patch_size = patch_size | |
| self.embed_dim = embed_dim | |
| self.Hp, self.Wp = H // patch_size, W // patch_size | |
| self.num_patches = self.Hp * self.Wp | |
| # split embed dim into thirds for spatial x, spatial y, and temporal | |
| base_split = (embed_dim // 3) & ~1 | |
| remaining_dim = embed_dim - base_split | |
| self.spatial_x_dim = (remaining_dim // 2) & ~1 | |
| self.spatial_y_dim = remaining_dim - self.spatial_x_dim | |
| self.temporal_dim = base_split | |
| # ensure the embed dim is split wholy into thirds and each third is even | |
| assert (self.spatial_x_dim + self.spatial_y_dim + self.temporal_dim) == embed_dim, \ | |
| f"Dimension mismatch: {self.spatial_x_dim} + {self.spatial_y_dim} + {self.temporal_dim} != {embed_dim}" | |
| assert self.spatial_x_dim % 2 == 0 and self.spatial_y_dim % 2 == 0 and self.temporal_dim % 2 == 0, \ | |
| f"Embed dim x={self.spatial_x_dim}, y={self.spatial_y_dim}, t={self.temporal_dim}" | |
| pe_spatial = build_spatial_only_pe(self.frame_size, self.patch_size, self.embed_dim, device='cpu', dtype=torch.float32) # [1,P,E] | |
| self.register_buffer("pos_spatial", pe_spatial, persistent=False) | |
| # pixel patches to embeddings | |
| self.proj = nn.Conv2d(3 * self.patch_size * self.patch_size, self.embed_dim, 1) | |
| def forward(self, frames): | |
| B, T, C, H, W = frames.shape | |
| # go from frames to patches | |
| x = rearrange(frames, 'b t c (hp p1) (wp p2) -> (b t) (c p1 p2) hp wp', p1=self.patch_size, p2=self.patch_size) # [(B*T), 3*p*p, Hp, Wp] | |
| x = self.proj(x) # [(B*T), E, Hp, Wp] | |
| x = rearrange(x, '(b t) e hp wp -> b t (hp wp) e', b=B, t=T) # [B, T, P, E] | |
| # add 2d spatial pos encoding (first 2/3 of embed dim) | |
| x = x + self.pos_spatial.to(dtype=x.dtype, device=x.device) # [B, T, P, E] | |
| return x | |