[CLEANUP]

zeta/nn/modules/patch_embedding_layer.py

@@ -0,0 +1,65 @@
+from torch import nn, Tensor
+from zeta.nn.modules.patch_img import patch_img
+from zeta.nn.attention.cross_attention import CrossAttention
+
+# from zeta.nn.modules.feedforward import Feedforward
+
+
+class PatchEmbeddingLayer(nn.Module):
+    def __init__(
+        self,
+        dim: int = None,
+        patches: int = 16,
+        image_size: int = 224,
+        in_channels: int = 3,
+    ):
+        super(PatchEmbeddingLayer, self).__init__()
+        self.dim = dim
+        self.patches = patches
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.patch_dim = in_channels * patches**2
+        self.patch_size = image_size // patches
+        self.num_patches = (image_size // self.patch_size) ** 2
+
+        self.cross_attn = CrossAttention(dim=dim, context_dim=self.dim)
+        self.ffn = nn.Sequential(
+            nn.Dropout(0.1),
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim),
+            nn.Linear(dim, dim * 4),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        patches = patch_img(
+            x,
+            patches=self.patches,
+        )
+        print(patches.shape)
+        b, s, d = patches.shape
+
+        # Run cross attn
+        # attended = self.cross_attn(patches, patches)
+        attended = CrossAttention(dim=d, context_dim=self.dim)(patches, patches)
+        print(attended.shape)
+
+        # Flatten patches
+        out = self.ffn(attended)
+        print(out.shape)
+
+        return out
+
+
+# x = torch.randn(1, 3, 224, 224)
+
+# model = PatchEmbeddingLayer(
+#     dim = 224,
+#     patches = 16,
+#     image_size = 224,
+#     in_channels = 3
+# )
+
+# out = model(x)
+# print(out.shape)