fal/moondream2-docci-instruct

1import torch2import torch.nn.functional as F3from torch import nn4from einops import rearrange5from torchvision.transforms.v2 import (6    Compose,7    Resize,8    InterpolationMode,9    ToImage,10    ToDtype,11    Normalize,12)13 14 15class Attention(nn.Module):16    def __init__(self, dim, num_heads=16):17        super().__init__()18        assert dim % num_heads == 0, "dim should be divisible by num_heads"19 20        self.num_heads = num_heads21        self.head_dim = dim // num_heads22 23        self.qkv = nn.Linear(dim, dim * 3)24        self.proj = nn.Linear(dim, dim)25 26        torch.nn.init.kaiming_normal_(27            self.qkv.weight, mode="fan_in", nonlinearity="relu"28        )29        torch.nn.init.kaiming_normal_(30            self.proj.weight, mode="fan_in", nonlinearity="relu"31        )32 33    def forward(self, x: torch.Tensor) -> torch.Tensor:34        B, N, C = x.shape35        qkv = (36            self.qkv(x)37            .reshape(B, N, 3, self.num_heads, self.head_dim)38            .permute(2, 0, 3, 1, 4)39        )40        q, k, v = qkv.unbind(0)41 42        x = F.scaled_dot_product_attention(q, k, v)43 44        x = x.transpose(1, 2).reshape(B, N, C)45        x = self.proj(x)46        return x47 48 49class VitBlock(nn.Module):50    def __init__(self, embed_dim):51        super().__init__()52        self.attn = Attention(embed_dim)53        self.mlp = MLP(embed_dim, 4304)54        self.norm1 = nn.LayerNorm(embed_dim)55        self.norm2 = nn.LayerNorm(embed_dim)56 57    def forward(self, x):58        x = x + self.attn(self.norm1(x))59        x = x + self.mlp(self.norm2(x))60        return x61 62 63class VisionTransformer(nn.Module):64 65    def __init__(self):66        super().__init__()67 68        embed_len = 72969        embed_dim = 115270 71        self.patch_embed = LinearPatchEmbedding()72        self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02)73        self.blocks = nn.Sequential(*[VitBlock(embed_dim) for _ in range(27)])74        self.norm = nn.LayerNorm(embed_dim)75 76    def forward(self, x):77        x = self.patch_embed(x)78        x = x + self.pos_embed79        for block in self.blocks:80            x = block(x)81        return self.norm(x)82 83 84class EncoderWrapper(nn.Module):85 86    def __init__(self):87        super().__init__()88        self.model = nn.ModuleDict({"visual": VisionTransformer()})89 90    def forward(self, x):91        return self.model["visual"](x)92 93 94class LinearPatchEmbedding(nn.Module):95 96    def __init__(self):97        super().__init__()98        self.linear = nn.Linear(588, 1152)99 100    def forward(self, x):101        return self.linear(x)102 103 104class MLP(nn.Module):105    def __init__(106        self,107        in_features: int,108        hidden_features: int = None,109        out_features: int = None,110    ) -> None:111        super().__init__()112        out_features = out_features or in_features113        hidden_features = hidden_features or in_features114        self.fc1 = nn.Linear(in_features, hidden_features)115        self.act = nn.GELU(approximate="tanh")116        self.fc2 = nn.Linear(hidden_features, out_features)117 118        torch.nn.init.kaiming_normal_(119            self.fc1.weight, mode="fan_in", nonlinearity="relu"120        )121        torch.nn.init.kaiming_normal_(122            self.fc2.weight, mode="fan_in", nonlinearity="relu"123        )124 125    def forward(self, x: torch.Tensor) -> torch.Tensor:126        x = self.fc1(x)127        x = self.act(x)128        x = self.fc2(x)129        return x130 131 132class VisionProjection(nn.Module):133    def __init__(self):134        super().__init__()135 136        image_embedding_dim = 1152137        model_dim = 2048138        hidden_dim = model_dim * 4139 140        self.mlp = MLP(image_embedding_dim, hidden_dim, model_dim)141 142    @property143    def device(self):144        return self.mlp.fc1.weight.device145 146    def forward(self, x):147        return self.mlp(x)148 149 150class VisionEncoder(nn.Module):151    def __init__(self) -> None:152        super().__init__()153 154        self.encoder = EncoderWrapper()155        self.projection = VisionProjection()156 157        self.preprocess = Compose(158            [159                Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC),160                ToImage(),161                ToDtype(torch.float32, scale=True),162                Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),163            ]164        )165 166    @property167    def device(self):168        return self.projection.mlp.fc1.weight.device169 170    @property171    def dtype(self):172        return self.projection.mlp.fc1.weight.dtype173 174    def __call__(self, images) -> torch.Tensor:175        if not isinstance(images, list):176            images = [images]177 178        with torch.no_grad():179            x = torch.stack(180                [self.preprocess(image.convert("RGB")) for image in images]181            ).to(self.device, dtype=self.dtype)182 183            x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14)184 185            x = self.encoder(x)186            x = self.projection(x)187 188            return x189