|
|
| from __future__ import annotations |
| import torch |
| import torch.nn as nn |
|
|
| class SquaredReLU(nn.Module): |
|
|
| def forward(self, x): |
| return torch.pow(torch.nn.functional.relu(x), 2) |
|
|
| class RMSNorm(nn.Module): |
|
|
| def __init__(self, hidden_size, eps=1e-05): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, hidden_states): |
| input_dtype = hidden_states.dtype |
| hidden_states = hidden_states.to(torch.float32) |
| variance = hidden_states.pow(2).mean(-1, keepdim=True) |
| hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| return (self.weight.to(torch.float32) * hidden_states).to(input_dtype) |
|
|
| class RadioHProjector(nn.Module): |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.downsample_ratio = 0.5 |
| self.ps_version = 'v2' |
| vit_hidden_size = 1280 |
| vision_projection_hidden_size = getattr(config, 'projector_hidden_size', config.hidden_size) |
| llm_hidden_size = config.hidden_size |
| in_dim = int(vit_hidden_size * (1 / self.downsample_ratio) ** 2) |
| self.mlp = nn.Sequential(RMSNorm(in_dim, eps=1e-05), nn.Linear(in_dim, vision_projection_hidden_size, bias=False), SquaredReLU(), nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False)) |
| self.radio_multiply = True |
|
|
| def pixel_shuffle(self, x, scale_factor=0.5): |
| (n, w, h, c) = x.size() |
| x = x.view(n, w, int(h * scale_factor), int(c / scale_factor)) |
| x = x.permute(0, 2, 1, 3).contiguous() |
| x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor))) |
| if self.ps_version != 'v1': |
| x = x.permute(0, 2, 1, 3).contiguous() |
| return x |
|
|
| def forward(self, image_features: torch.Tensor, image_grid_thw: torch.Tensor): |
| device = image_features.device |
| indices = [] |
| offset = 0 |
| for (_, H, W) in image_grid_thw.tolist(): |
| (H, W) = (int(H), int(W)) |
| (H2, W2) = (H // 2, W // 2) |
| k = torch.arange(H2 * W2, device=device) |
| if self.ps_version == 'v1': |
| (out_r, out_c) = (k % H2, k // H2) |
| else: |
| (out_r, out_c) = (k // W2, k % W2) |
| tl = offset + 2 * out_r * W + 2 * out_c |
| indices.append(torch.stack([tl, tl + 1, tl + W, tl + W + 1], dim=1).reshape(-1)) |
| offset += H * W |
| gather_idx = torch.cat(indices) |
| merged = image_features[gather_idx].view(-1, 4 * image_features.shape[-1]) |
| if self.radio_multiply: |
| return 0.01 * self.mlp(merged) |
| return self.mlp(merged) |
|
|
