| import torch |
| import torch.nn as nn |
| import math |
|
|
| from .layers import linear, mlp |
|
|
|
|
| def fourier_features(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: |
| """ |
| Applies Fourier feature mapping to input tensor x using frequency matrix w. This |
| projects inputs through sinusoidal functions to create higher dimensional features |
| that help mitigate spectral bias - the tendency of neural networks to learn |
| low-frequency functions more easily than high-frequency ones. By explicitly |
| mapping inputs to higher frequencies through sin/cos transformations, we enable |
| better learning of fine details and higher frequency patterns. |
| |
| Args: |
| x: Input tensor to transform |
| w: Matrix of frequencies for the Fourier features transformation |
| |
| Returns: |
| Concatenated cosine and sine transformed features as a tensor |
| """ |
| f = 2 * math.pi * x @ w |
| return torch.cat([f.cos(), f.sin()], dim=-1) |
|
|
|
|
| def encode_coordinate(coord: torch.Tensor, w: nn.Module) -> torch.Tensor: |
| """ |
| Takes as input a tensor containing a single float coordinate value (x or y) |
| and encodes it into hidden states for input to the text model. |
| |
| Args: |
| coord: Tensor with single float coordinate value |
| |
| Returns: |
| Encoded hidden states tensor for input to text model |
| """ |
| return linear(fourier_features(coord, w.coord_features), w.coord_encoder) |
|
|
|
|
| def decode_coordinate(hidden_state: torch.Tensor, w: nn.Module) -> torch.Tensor: |
| """ |
| Takes as input the last hidden state from the text model and outputs a single logit |
| representing either an x or y coordinate prediction. |
| |
| Args: |
| hidden_state: The final hidden state tensor from the text model. |
| |
| Returns: |
| A single logit representing the predicted coordinate value (x or y) |
| """ |
| return mlp(hidden_state, w.coord_decoder) |
|
|
|
|
| def encode_size(size: torch.Tensor, w: nn.Module) -> torch.Tensor: |
| """ |
| Takes a tensor containing width and height values and encodes them into |
| hidden states for input to the text model. |
| |
| Args: |
| size: Tensor with two floats for width and height |
| |
| Returns: |
| Encoded hidden states tensor for input to text model |
| """ |
| return linear(fourier_features(size, w.size_features), w.size_encoder) |
|
|
|
|
| def decode_size(hidden_state: torch.Tensor, w: nn.Module) -> torch.Tensor: |
| """ |
| Takes as input the last hidden state from the text model and outputs logits |
| for 1024 bins representing width and height in log-scale. |
| |
| The bins are distributed according to the formula: |
| bin = (log2(size) + 10.0) / 10.0 * 1023.0 |
| where size values are clamped to be at least 1/1024. |
| |
| To convert from bin back to size: |
| size = 2^((bin / 1023.0) * 10.0 - 10.0) |
| |
| Args: |
| hidden_state: The final hidden state tensor from the text model. |
| |
| Returns: |
| A tensor containing logits for 1024 bins for width and height. |
| Shape is (2, 1024) where the first dimension corresponds to width and height. |
| """ |
| return mlp(hidden_state, w.size_decoder).view(2, -1) |
|
|
