Update vision_encoder.py

vision_encoder.py

@@ -26,7 +26,6 @@ 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"
@@ -76,11 +75,10 @@ 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, 4304)
+        self.mlp = MLP(embed_dim, 2304)
         self.norm1 = nn.LayerNorm(embed_dim)
         self.norm2 = nn.LayerNorm(embed_dim)
 
@@ -91,14 +89,13 @@ class VitBlock(nn.Module):
 
 
 class VisionTransformer(nn.Module):
-
     def __init__(self, use_flash_attn=False):
         super().__init__()
 
         embed_len = 729
-        embed_dim = 4608
+        embed_dim = 1152  # Updated to match checkpoint
 
-        self.patch_embed = LinearPatchEmbedding()
+        self.patch_embed = LinearPatchEmbedding(embed_dim)
         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)]
@@ -113,21 +110,10 @@ 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):
+    def __init__(self, embed_dim=1152):  # Updated default to match checkpoint
         super().__init__()
-        self.linear = nn.Linear(588, 4608)
+        self.linear = nn.Linear(588, embed_dim)  # Match saved model
 
     def forward(self, x):
         b, c, hp1, wp2 = x.shape
@@ -136,190 +122,4 @@ 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 = 2304
-        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, 4608, 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(4608, 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, 4608, 729).transpose(1, 2)
-        )
-
-        final_features = torch.cat([full_img_features, reshaped_patch_features], dim=2)
-
-        return self.projection(final_features)