| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> torch.Tensor:
|
| grid_h = torch.arange(grid_size, dtype=torch.float32)
|
| grid_w = torch.arange(grid_size, dtype=torch.float32)
|
| grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
|
| grid = torch.stack(grid, dim=0).reshape([2, 1, grid_size, grid_size])
|
|
|
| dim = embed_dim // 2
|
| omega = torch.arange(dim // 2, dtype=torch.float32) / (dim / 2.0)
|
| omega = 1.0 / (10000 ** omega)
|
|
|
| out_w = torch.einsum("m,d->md", grid[0].reshape(-1), omega)
|
| out_h = torch.einsum("m,d->md", grid[1].reshape(-1), omega)
|
|
|
| pos_embed = torch.cat([
|
| torch.cat([torch.sin(out_w), torch.cos(out_w)], dim=1),
|
| torch.cat([torch.sin(out_h), torch.cos(out_h)], dim=1)
|
| ], dim=1)
|
|
|
| if cls_token:
|
| pos_embed = torch.cat([torch.zeros([1, embed_dim], dtype=torch.float32), pos_embed], dim=0)
|
| return pos_embed
|
|
|
|
|
| class ConvNeXtStem(nn.Module):
|
| """Sub-pixel ConvNeXt Stem preserving high-frequency edge gradients."""
|
| def __init__(self, in_channels: int = 3, embed_dim: int = 256):
|
| super().__init__()
|
| mid_dim = max(64, embed_dim // 2)
|
| self.stem = nn.Sequential(
|
| nn.Conv2d(in_channels, mid_dim, kernel_size=7, stride=2, padding=3),
|
| nn.BatchNorm2d(mid_dim),
|
| nn.GELU(),
|
| nn.Conv2d(mid_dim, embed_dim, kernel_size=3, stride=2, padding=1),
|
| nn.BatchNorm2d(embed_dim),
|
| nn.GELU(),
|
| nn.Conv2d(embed_dim, embed_dim, kernel_size=4, stride=4),
|
| nn.BatchNorm2d(embed_dim),
|
| )
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| feat = self.stem(x)
|
| return feat.flatten(2).transpose(1, 2)
|
|
|
|
|
| class TransformerBlock(nn.Module):
|
| def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0):
|
| super().__init__()
|
| self.norm1 = nn.LayerNorm(dim, eps=1e-6)
|
| self.num_heads = num_heads
|
| self.head_dim = dim // num_heads
|
| self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
| self.proj = nn.Linear(dim, dim)
|
| self.norm2 = nn.LayerNorm(dim, eps=1e-6)
|
| mlp_hidden_dim = int(dim * mlp_ratio)
|
| self.mlp = nn.Sequential(
|
| nn.Linear(dim, mlp_hidden_dim),
|
| nn.GELU(),
|
| nn.Linear(mlp_hidden_dim, dim),
|
| )
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| B, N, C = x.shape
|
| x_norm = self.norm1(x)
|
| qkv = self.qkv(x_norm).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
|
| q, k, v = qkv[0], qkv[1], qkv[2]
|
| attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
|
| attn_out = attn_out.transpose(1, 2).reshape(B, N, C)
|
| x = x + self.proj(attn_out)
|
| x = x + self.mlp(self.norm2(x))
|
| return x
|
|
|
|
|
| class AevanisVisionBackbone(nn.Module):
|
| """
|
| Aevanis Vision Foundation Backbone (Universal Visual Eye)
|
| - Nano: 7.68M parameters (embed_dim=256, depth=8, heads=4)
|
| - Core: 24.35M parameters (embed_dim=384, depth=12, heads=6)
|
| """
|
| def __init__(self, model_scale: str = "nano", img_size: int = 224, patch_size: int = 16, in_channels: int = 3):
|
| super().__init__()
|
| self.model_scale = model_scale
|
| if model_scale == "nano":
|
| embed_dim, depth, num_heads = 256, 8, 4
|
| elif model_scale == "core":
|
| embed_dim, depth, num_heads = 384, 12, 6
|
| else:
|
| raise ValueError(f"Unknown scale: {model_scale}")
|
|
|
| self.embed_dim = embed_dim
|
| self.patch_embed = ConvNeXtStem(in_channels=in_channels, embed_dim=embed_dim)
|
| grid_size = img_size // patch_size
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
| pos_embed = get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=True)
|
| self.register_buffer("pos_embed", pos_embed.unsqueeze(0))
|
|
|
| self.blocks = nn.ModuleList([
|
| TransformerBlock(embed_dim, num_heads, mlp_ratio=4.0) for _ in range(depth)
|
| ])
|
| self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
|
|
|
| def forward(self, x: torch.Tensor):
|
| B = x.shape[0]
|
| x = self.patch_embed(x)
|
| x = x + self.pos_embed[:, 1:, :]
|
| cls_tokens = self.cls_token.expand(B, -1, -1) + self.pos_embed[:, :1, :]
|
| x = torch.cat((cls_tokens, x), dim=1)
|
| for blk in self.blocks:
|
| x = blk(x)
|
| x = self.norm(x)
|
| cls_token = x[:, 0]
|
| spatial_tokens = x[:, 1:]
|
| return {
|
| "cls_token": cls_token,
|
| "spatial_tokens": spatial_tokens,
|
| "full_tokens": x
|
| }
|
|
|
| def forward_encoder(self, x: torch.Tensor, mask_ratio: float = 0.0):
|
|
|
| out = self.forward(x)
|
| return out["full_tokens"], None, None
|
|
|
| @classmethod
|
| def from_pretrained(cls, weight_path: str, model_scale: str = "nano"):
|
| model = cls(model_scale=model_scale)
|
| ckpt = torch.load(weight_path, map_location="cpu")
|
| cleaned = {k.replace("module.", ""): v for k, v in ckpt.items() if not k.startswith("decoder")}
|
| model.load_state_dict(cleaned, strict=False)
|
| model.eval()
|
| return model
|
|
|
