| """Shared image model components.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ImageEncoder(nn.Module): |
| def __init__(self, emb_dim: int = 96): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Conv2d(5, 24, kernel_size=5, stride=2, padding=2), |
| nn.SiLU(), |
| nn.Conv2d(24, 48, kernel_size=3, stride=2, padding=1), |
| nn.SiLU(), |
| nn.Conv2d(48, 64, kernel_size=3, stride=2, padding=1), |
| nn.SiLU(), |
| nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), |
| nn.SiLU(), |
| nn.AdaptiveAvgPool2d((8, 8)), |
| nn.Flatten(), |
| nn.Linear(64 * 8 * 8, emb_dim), |
| nn.LayerNorm(emb_dim), |
| ) |
|
|
| def forward(self, images: torch.Tensor) -> torch.Tensor: |
| x = images.float() / 255.0 |
| if x.is_cuda: |
| x = x.contiguous(memory_format=torch.channels_last) |
| b, _c, h, w = x.shape |
| yy = torch.linspace(-1.0, 1.0, h, device=x.device, dtype=x.dtype).view(1, 1, h, 1).expand(b, 1, h, w) |
| xx = torch.linspace(-1.0, 1.0, w, device=x.device, dtype=x.dtype).view(1, 1, 1, w).expand(b, 1, h, w) |
| x = torch.cat([x, xx, yy], dim=1) |
| return self.net(x) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, in_dim: int, out_dim: int, hidden_dim: int = 160, depth: int = 2): |
| super().__init__() |
| layers: list[nn.Module] = [] |
| dim = in_dim |
| for _ in range(depth): |
| layers.append(nn.Linear(dim, hidden_dim)) |
| layers.append(nn.SiLU()) |
| dim = hidden_dim |
| layers.append(nn.Linear(dim, out_dim)) |
| self.net = nn.Sequential(*layers) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.net(x) |
|
|
|
|
| def encode_image_sequence(encoder: ImageEncoder, images: torch.Tensor) -> torch.Tensor: |
| b, t, c, h, w = images.shape |
| emb = encoder(images.reshape(b * t, c, h, w)) |
| return emb.reshape(b, t, -1) |
|
|
