|
import torch |
|
import torch.nn.functional as F |
|
from torch import nn |
|
from einops import rearrange |
|
from torchvision.transforms.v2 import ( |
|
Compose, |
|
Resize, |
|
InterpolationMode, |
|
ToImage, |
|
ToDtype, |
|
Normalize, |
|
) |
|
from transformers.utils import is_flash_attn_2_available |
|
|
|
try: |
|
if is_flash_attn_2_available(): |
|
from flash_attn.modules.mha import FlashSelfAttention |
|
else: |
|
FlashSelfAttention = None |
|
except ImportError: |
|
FlashSelfAttention = None |
|
|
|
|
|
class Attention(nn.Module): |
|
|
|
def __init__(self, dim, num_heads=16, use_flash_attn=False): |
|
super().__init__() |
|
assert dim % num_heads == 0, "dim should be divisible by num_heads" |
|
|
|
self.num_heads = num_heads |
|
self.head_dim = dim // num_heads |
|
|
|
self.qkv = nn.Linear(dim, dim * 3) |
|
self.proj = nn.Linear(dim, dim) |
|
|
|
if use_flash_attn and FlashSelfAttention is not None: |
|
self.flash_attn = FlashSelfAttention() |
|
else: |
|
self.flash_attn = None |
|
|
|
torch.nn.init.kaiming_normal_( |
|
self.qkv.weight, mode="fan_in", nonlinearity="relu" |
|
) |
|
torch.nn.init.kaiming_normal_( |
|
self.proj.weight, mode="fan_in", nonlinearity="relu" |
|
) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
if self.flash_attn is not None: |
|
qkv = self.qkv(x) |
|
qkv = rearrange( |
|
qkv, "... (three h d) -> ... three h d", three=3, h=self.num_heads |
|
) |
|
attn_output = self.flash_attn(qkv) |
|
output = rearrange(attn_output, "... h d -> ... (h d)") |
|
output = self.proj(output) |
|
return output |
|
else: |
|
B, N, C = x.shape |
|
qkv = ( |
|
self.qkv(x) |
|
.reshape(B, N, 3, self.num_heads, self.head_dim) |
|
.permute(2, 0, 3, 1, 4) |
|
) |
|
q, k, v = qkv.unbind(0) |
|
|
|
x = F.scaled_dot_product_attention(q, k, v) |
|
|
|
x = x.transpose(1, 2).reshape(B, N, C) |
|
x = self.proj(x) |
|
return x |
|
|
|
|
|
class VitBlock(nn.Module): |
|
|
|
def __init__(self, embed_dim, use_flash_attn=False): |
|
super().__init__() |
|
self.attn = Attention(embed_dim, use_flash_attn=use_flash_attn) |
|
self.mlp = MLP(embed_dim, 4304) |
|
self.norm1 = nn.LayerNorm(embed_dim) |
|
self.norm2 = nn.LayerNorm(embed_dim) |
|
|
|
def forward(self, x): |
|
x = x + self.attn(self.norm1(x)) |
|
x = x + self.mlp(self.norm2(x)) |
|
return x |
|
|
|
|
|
class VisionTransformer(nn.Module): |
|
|
|
def __init__(self, use_flash_attn=False): |
|
super().__init__() |
|
|
|
embed_len = 729 |
|
embed_dim = 1152 |
|
|
|
self.patch_embed = LinearPatchEmbedding() |
|
self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02) |
|
self.blocks = nn.Sequential( |
|
*[VitBlock(embed_dim, use_flash_attn=use_flash_attn) for _ in range(27)] |
|
) |
|
self.norm = nn.LayerNorm(embed_dim) |
|
|
|
def forward(self, x): |
|
x = self.patch_embed(x) |
|
x = x + self.pos_embed |
|
for block in self.blocks: |
|
x = block(x) |
|
return self.norm(x) |
|
|
|
|
|
class EncoderWrapper(nn.Module): |
|
|
|
def __init__(self, use_flash_attn=False): |
|
super().__init__() |
|
self.model = nn.ModuleDict({"visual": VisionTransformer(use_flash_attn)}) |
|
|
|
def forward(self, x): |
|
return self.model["visual"](x) |
|
|
|
|
|
class LinearPatchEmbedding(nn.Module): |
|
|
|
def __init__(self): |
|
super().__init__() |
|
self.linear = nn.Linear(588, 1152) |
|
|
|
def forward(self, x): |
|
b, c, hp1, wp2 = x.shape |
|
p1, p2 = 14, 14 |
|
h, w = hp1 // p1, wp2 // p2 |
|
x = x.reshape(b, c, h, p1, w, p2) |
|
x = x.permute(0, 2, 4, 1, 3, 5) |
|
x = x.reshape(b, h * w, c * p1 * p2) |
|
|
|
return self.linear(x) |
|
|
|
|
|
class MLP(nn.Module): |
|
def __init__( |
|
self, |
|
in_features: int, |
|
hidden_features: int = None, |
|
out_features: int = None, |
|
) -> None: |
|
super().__init__() |
|
out_features = out_features or in_features |
|
hidden_features = hidden_features or in_features |
|
self.fc1 = nn.Linear(in_features, hidden_features) |
|
self.act = nn.GELU(approximate="tanh") |
|
self.fc2 = nn.Linear(hidden_features, out_features) |
|
|
|
torch.nn.init.kaiming_normal_( |
|
self.fc1.weight, mode="fan_in", nonlinearity="relu" |
|
) |
|
torch.nn.init.kaiming_normal_( |
|
self.fc2.weight, mode="fan_in", nonlinearity="relu" |
|
) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
x = self.fc1(x) |
|
x = self.act(x) |
|
x = self.fc2(x) |
|
return x |
|
|
|
|
|
class VisionProjection(nn.Module): |
|
def __init__(self): |
|
super().__init__() |
|
|
|
image_embedding_dim = 1152 |
|
model_dim = 2048 |
|
hidden_dim = model_dim * 4 |
|
|
|
self.mlp = MLP(image_embedding_dim, hidden_dim, model_dim) |
|
|
|
@property |
|
def device(self): |
|
return self.mlp.fc1.weight.device |
|
|
|
def forward(self, x): |
|
return self.mlp(x) |
|
|
|
|
|
class VisionEncoder(nn.Module): |
|
|
|
def __init__(self, use_flash_attn=False): |
|
super().__init__() |
|
|
|
self.encoder = EncoderWrapper(use_flash_attn) |
|
self.projection = VisionProjection() |
|
|
|
self.preprocess = Compose( |
|
[ |
|
Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC), |
|
ToImage(), |
|
ToDtype(torch.float32, scale=True), |
|
Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), |
|
] |
|
) |
|
|
|
@property |
|
def device(self): |
|
return self.projection.mlp.fc1.weight.device |
|
|
|
@property |
|
def dtype(self): |
|
return self.projection.mlp.fc1.weight.dtype |
|
|
|
def __call__(self, images) -> torch.Tensor: |
|
if not isinstance(images, list) and not isinstance(images, torch.Tensor): |
|
images = [images] |
|
|
|
with torch.no_grad(): |
|
|
|
if not isinstance(images, torch.Tensor) and not isinstance( |
|
images[0], torch.Tensor |
|
): |
|
images = [self.preprocess(image.convert("RGB")) for image in images] |
|
|
|
if isinstance(images, list): |
|
images = torch.stack(images) |
|
|
|
x = images.to(self.device, dtype=self.dtype) |
|
x = self.encoder(x) |
|
x = self.projection(x) |
|
|
|
return x |
|
|
