1 files changed, 90 insertions, 0 deletions
diff --git a/worldalign/models.py b/worldalign/models.py
new file mode 100644
index 0000000..5f98605
--- /dev/null
+++ b/worldalign/models.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class Bridge(nn.Module):
+    def __init__(
+        self,
+        vision_dim: int,
+        text_dim: int,
+        hidden_dim: int = 1536,
+        linear: bool = False,
+    ):
+        super().__init__()
+        self.vision_dim = vision_dim
+        self.text_dim = text_dim
+        self.hidden_dim = hidden_dim
+        self.linear = linear
+        if linear:
+            self.network = nn.Linear(vision_dim, text_dim, bias=False)
+        else:
+            self.network = nn.Sequential(
+                nn.LayerNorm(vision_dim),
+                nn.Linear(vision_dim, hidden_dim),
+                nn.GELU(),
+                nn.Linear(hidden_dim, text_dim),
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.normalize(self.network(x).float(), dim=-1)
+
+    def config(self) -> dict:
+        return {
+            "vision_dim": self.vision_dim,
+            "text_dim": self.text_dim,
+            "hidden_dim": self.hidden_dim,
+            "linear": self.linear,
+        }
+
+
+class PrefixAdapter(nn.Module):
+    def __init__(
+        self,
+        semantic_dim: int,
+        lm_dim: int,
+        prefix_length: int = 8,
+        hidden_dim: int = 2048,
+    ):
+        super().__init__()
+        self.semantic_dim = semantic_dim
+        self.lm_dim = lm_dim
+        self.prefix_length = prefix_length
+        self.hidden_dim = hidden_dim
+        self.network = nn.Sequential(
+            nn.LayerNorm(semantic_dim),
+            nn.Linear(semantic_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, prefix_length * lm_dim),
+        )
+
+    def forward(self, semantic: torch.Tensor) -> torch.Tensor:
+        prefix = self.network(semantic.float())
+        return prefix.view(-1, self.prefix_length, self.lm_dim)
+
+    def config(self) -> dict:
+        return {
+            "semantic_dim": self.semantic_dim,
+            "lm_dim": self.lm_dim,
+            "prefix_length": self.prefix_length,
+            "hidden_dim": self.hidden_dim,
+        }
+
+
+def load_bridge(path: str, device: str = "cpu") -> tuple[Bridge, dict]:
+    state = torch.load(path, map_location="cpu", weights_only=False)
+    model = Bridge(**state["config"])
+    model.load_state_dict(state["state_dict"])
+    model.to(device).eval()
+    return model, state
+
+
+def load_prefix(path: str, device: str = "cpu") -> tuple[PrefixAdapter, dict]:
+    state = torch.load(path, map_location="cpu", weights_only=False)
+    model = PrefixAdapter(**state["config"])
+    model.load_state_dict(state["state_dict"])
+    model.to(device).eval()
+    return model, state
+