Upload vision_encoder.py

vision_encoder.py

@@ -26,6 +26,7 @@ except ImportError:
 
 
 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"
@@ -75,10 +76,11 @@ class Attention(nn.Module):
 
 
 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, 2304)
+        self.mlp = MLP(embed_dim, 4304)
         self.norm1 = nn.LayerNorm(embed_dim)
         self.norm2 = nn.LayerNorm(embed_dim)
 
@@ -89,13 +91,14 @@ class VitBlock(nn.Module):
 
 
 class VisionTransformer(nn.Module):
+
     def __init__(self, use_flash_attn=False):
         super().__init__()
 
         embed_len = 729
-        embed_dim = 1152  # Updated to match checkpoint
+        embed_dim = 1152
 
-        self.patch_embed = LinearPatchEmbedding(embed_dim)
+        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)]
@@ -110,10 +113,21 @@ class VisionTransformer(nn.Module):
         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, embed_dim=1152):  # Updated default to match checkpoint
+
+    def __init__(self):
         super().__init__()
-        self.linear = nn.Linear(588, embed_dim)  # Match saved model
+        self.linear = nn.Linear(588, 1152)
 
     def forward(self, x):
         b, c, hp1, wp2 = x.shape
@@ -122,4 +136,190 @@ class LinearPatchEmbedding(nn.Module):
         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 * 2, hidden_dim, model_dim)
+
+    @property
+    def device(self):
+        return self.mlp.fc1.weight.device
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+def create_patches(image, patch_size=(378, 378)):
+    assert image.dim() == 3, "Image must be in CHW format"
+
+    _, height, width = image.shape  # Channels, Height, Width
+    patch_height, patch_width = patch_size
+
+    if height == patch_height and width == patch_width:
+        return []
+
+    # Iterate over the image and create patches
+    patches = []
+    for i in range(0, height, patch_height):
+        row_patches = []
+        for j in range(0, width, patch_width):
+            patch = image[:, i : i + patch_height, j : j + patch_width]
+            row_patches.append(patch)
+        patches.append(torch.stack(row_patches))
+    return patches
+
+
+class VisionEncoder(nn.Module):
+
+    def __init__(self, use_flash_attn=False):
+        super().__init__()
+
+        self.encoder = EncoderWrapper(use_flash_attn)
+        self.projection = VisionProjection()
+        self.supported_sizes = [(378, 378), (378, 756), (756, 378), (756, 756)]
+
+    @property
+    def device(self):
+        return self.projection.mlp.fc1.weight.device
+
+    @property
+    def dtype(self):
+        return self.projection.mlp.fc1.weight.dtype
+
+    def preprocess(self, image: PIL.Image.Image):
+        width, height = image.size
+        max_dim = max(width, height)
+        if max_dim < 512:
+            im_size = (378, 378)
+        else:
+            aspect_ratio = width / height
+            im_size = min(
+                self.supported_sizes,
+                key=lambda size: (
+                    abs((size[1] / size[0]) - aspect_ratio),
+                    abs(size[0] - width) + abs(size[1] - height),
+                ),
+            )
+
+        return Compose(
+            [
+                Resize(size=im_size, interpolation=InterpolationMode.BICUBIC),
+                ToImage(),
+                ToDtype(torch.float32, scale=True),
+                Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+            ]
+        )(image)
+
+    def forward(
+        self, images: Union[PIL.Image.Image, list[PIL.Image.Image], torch.Tensor]
+    ) -> torch.Tensor:
+        im_list = None
+        if isinstance(images, torch.Tensor):
+            # Input must have dimensions (B, C, H, W)
+            assert (
+                len(images.shape) == 4
+            ), "Tensor input must have dimensions (B, C, H, W)"
+            im_list = list(images)
+        elif isinstance(images, PIL.Image.Image):
+            im_list = [images]
+        elif isinstance(images, list):
+            im_list = images
+        else:
+            raise ValueError(
+                "Input must be a PIL image, list of PIL images, or a tensor"
+            )
+
+        # Preprocess unless the images are already tensors (indicating that
+        # they have already been preprocessed)
+        if not isinstance(im_list[0], torch.Tensor):
+            im_list = [self.preprocess(im.convert("RGB")) for im in im_list]
+
+        patches = [create_patches(im) for im in im_list]
+        flat_patches = [patch for image_patches in patches for patch in image_patches]
+
+        # Images may be variable size, and need to be resized to a common size after
+        # creating patches.
+        resized_images = [
+            F.interpolate(im.unsqueeze(0), size=(378, 378), mode="bilinear")
+            for im in im_list
+        ]
+
+        combined_images = torch.cat([*resized_images, *flat_patches], dim=0)
+        combined_images = combined_images.to(self.device, dtype=self.dtype)
+
+        combined_features = self.encoder(combined_images)
+
+        full_img_features = combined_features[: len(im_list)]
+        patch_features = (
+            combined_features[len(im_list) :].transpose(1, 2).view(-1, 1152, 27, 27)
+        )
+
+        # Reshape patch features back to their original structure
+        reshaped_patch_features = []
+        patch_idx = 0
+        for i, patch_set in enumerate(patches):
+            if len(patch_set) == 0:
+                reshaped_patch_features.append(
+                    full_img_features[i].transpose(0, 1).view(1152, 27, 27)
+                )
+            else:
+                sample_features = []
+                for row_patches in patch_set:
+                    row_len = len(row_patches)
+                    row_features = patch_features[
+                        patch_idx : patch_idx + row_len
+                    ]  # row_len, T, C
+                    row_features = torch.cat(
+                        list(row_features), dim=2
+                    )  # T, C * row_len
+                    patch_idx += row_len
+                    sample_features.append(row_features)
+                sample_features = torch.cat(sample_features, dim=1)
+                sample_features = F.interpolate(
+                    sample_features.unsqueeze(0), size=(27, 27), mode="bilinear"
+                ).squeeze(0)
+                reshaped_patch_features.append(sample_features)
+        reshaped_patch_features = (
+            torch.stack(reshaped_patch_features).view(-1, 1152, 729).transpose(1, 2)
+        )
+
+        final_features = torch.cat([full_img_features, reshaped_patch_features], dim=2)
+
+        return self.projection(final_features)