Upload transformer/model.py with huggingface_hub

transformer/model.py

@@ -936,16 +936,14 @@ class DiffusionTokenEncoder(nn.Module):
 
 
 class EmbeddingLayer(nn.Module):
-    """Embedding layer for 1D features."""
+    """Embedding layer for 1D features - simple linear projection."""
 
-    def __init__(self, n_channels: int, total_channels: int, output_channels: int):
+    def __init__(self, n_channels: int, output_channels: int):
         super().__init__()
-        self.weight = nn.Parameter(torch.zeros(n_channels, total_channels))
-        self.proj = linearNoBias(total_channels, output_channels)
+        self.weight = nn.Parameter(torch.zeros(output_channels, n_channels))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        emb = torch.einsum("...i,io->...o", x, self.weight)
-        return self.proj(emb)
+        return F.linear(x, self.weight)
 
 
 class OneDFeatureEmbedder(nn.Module):
@@ -954,13 +952,12 @@ class OneDFeatureEmbedder(nn.Module):
     def __init__(self, features: dict, output_channels: int):
         super().__init__()
         self.features = {k: v for k, v in features.items() if v is not None}
-        total_embedding_input_features = sum(self.features.values())
         self.embedders = nn.ModuleDict({
-            feature: EmbeddingLayer(n_channels, total_embedding_input_features, output_channels)
+            feature: EmbeddingLayer(n_channels, output_channels)
             for feature, n_channels in self.features.items()
         })
 
-    def forward(self, f: dict, collapse_length: int) -> torch.Tensor:
+    def forward(self, f: dict) -> torch.Tensor:
         result = None
         for feature in self.features:
             x = f.get(feature)