|
|
| from __future__ import annotations |
| import torch |
| from torch import nn |
| IMAGE_LOCAL_AR_CHILD_COUNT = 4 |
|
|
| class PackedLinearImageCodec(nn.Module): |
|
|
| def __init__(self, channels: int, *, norm_eps: float=1e-06, device: torch.device | str | None=None, dtype: torch.dtype | None=None) -> None: |
| super().__init__() |
| if isinstance(channels, bool) or not isinstance(channels, int) or channels < 1: |
| raise ValueError('channels must be a positive integer') |
| factory_kwargs = {'device': device, 'dtype': dtype} |
| self.channels = channels |
| self.down_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs) |
| self.down_projection = nn.Linear(4 * channels, channels, bias=False, **factory_kwargs) |
| self.up_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs) |
| self.up_projection = nn.Linear(channels, 4 * channels, bias=False, **factory_kwargs) |
|
|
| @staticmethod |
| def _shape(shape: tuple[int, int] | list[int]) -> tuple[int, int]: |
| if not isinstance(shape, (tuple, list)) or len(shape) != 2: |
| raise ValueError('image token shape must contain height and width') |
| (height, width) = shape |
| if any((isinstance(value, bool) or not isinstance(value, int) for value in shape)): |
| raise TypeError('image token height and width must be integers') |
| if height < 1 or width < 1 or height % 2 or width % 2: |
| raise ValueError('packed-linear image grids must have positive even dimensions') |
| return (height, width) |
|
|
| @staticmethod |
| def _patches(grid: torch.Tensor) -> torch.Tensor: |
| return torch.stack((grid[..., 0::2, 0::2, :], grid[..., 0::2, 1::2, :], grid[..., 1::2, 0::2, :], grid[..., 1::2, 1::2, :]), dim=-2) |
|
|
| @classmethod |
| def pack_raster_values(cls, child_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor: |
| (height, width) = cls._shape(shape) |
| if not isinstance(child_values, torch.Tensor): |
| raise TypeError('child_values must be a tensor') |
| if child_values.ndim < 2 or child_values.shape[1] != height * width: |
| raise ValueError(f'child_values must have shape [B, {height * width}, ...]') |
| grid = child_values.reshape(child_values.shape[0], height, width, *child_values.shape[2:]) |
| patches = torch.stack((grid[:, 0::2, 0::2, ...], grid[:, 0::2, 1::2, ...], grid[:, 1::2, 0::2, ...], grid[:, 1::2, 1::2, ...]), dim=3) |
| return patches.flatten(1, 2) |
|
|
| @classmethod |
| def unpack_patch_values(cls, patch_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor: |
| (height, width) = cls._shape(shape) |
| (parent_height, parent_width) = (height // 2, width // 2) |
| parent_count = parent_height * parent_width |
| if not isinstance(patch_values, torch.Tensor): |
| raise TypeError('patch_values must be a tensor') |
| if patch_values.ndim < 3 or patch_values.shape[1] != parent_count or patch_values.shape[2] != IMAGE_LOCAL_AR_CHILD_COUNT: |
| raise ValueError(f'patch_values must have shape [B, {parent_count}, {IMAGE_LOCAL_AR_CHILD_COUNT}, ...]') |
| patches = patch_values.reshape(patch_values.shape[0], parent_height, parent_width, IMAGE_LOCAL_AR_CHILD_COUNT, *patch_values.shape[3:]) |
| grid = patch_values.new_empty(patch_values.shape[0], height, width, *patch_values.shape[3:]) |
| grid[:, 0::2, 0::2, ...] = patches[:, :, :, 0, ...] |
| grid[:, 0::2, 1::2, ...] = patches[:, :, :, 1, ...] |
| grid[:, 1::2, 0::2, ...] = patches[:, :, :, 2, ...] |
| grid[:, 1::2, 1::2, ...] = patches[:, :, :, 3, ...] |
| return grid.flatten(1, 2) |
|
|
| def compress_fixed_grid(self, child_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor: |
| (height, width) = self._shape(shape) |
| if child_states.ndim != 3 or child_states.shape[1:] != (height * width, self.channels): |
| raise ValueError(f'child_states must have shape [B, {height * width}, {self.channels}]') |
| grid = child_states.reshape(child_states.shape[0], height, width, self.channels) |
| patches = self._patches(grid) |
| normalized = self.down_norm(patches) |
| parents = self.down_projection(normalized.reshape(child_states.shape[0], height // 2, width // 2, 4 * self.channels)) |
| return parents.flatten(1, 2) |
|
|
| def expand_fixed_grid(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor: |
| children = self.expand_fixed_grid_patches(parent_states, shape) |
| return self.unpack_patch_values(children, shape) |
|
|
| def expand_fixed_grid_patches(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor: |
| (height, width) = self._shape(shape) |
| parent_count = height // 2 * (width // 2) |
| if parent_states.ndim != 3 or parent_states.shape[1:] != (parent_count, self.channels): |
| raise ValueError(f'parent_states must have shape [B, {parent_count}, {self.channels}]') |
| return self.up_projection(self.up_norm(parent_states)).reshape(parent_states.shape[0], parent_count, 4, self.channels) |
|
|
