Faz2 dep: code/meshai_train/models.py

code/meshai_train/models.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from meshai_bridge.latent_adapter import TrellisHunyuanLatentAdapter
+
+GEOM_IN_DIM = 4096 + 6
+VOXEL_OUT_DIM = 32 * 32 * 32
+TEXTURE_LATENT_DIM = 512
+
+
+class GeometryVoxelHead(nn.Module):
+    def __init__(self, in_dim: int = GEOM_IN_DIM, out_dim: int = VOXEL_OUT_DIM) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, 2048),
+            nn.GELU(),
+            nn.Linear(2048, 4096),
+            nn.GELU(),
+            nn.Linear(4096, out_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class TextureViewEncoder(nn.Module):
+    """Multi-view CNN -> conditioning vector for Hunyuan path."""
+
+    def __init__(self, out_dim: int = TEXTURE_LATENT_DIM, render_size: int = 128) -> None:
+        super().__init__()
+        self.render_size = render_size
+        self.backbone = nn.Sequential(
+            nn.Conv2d(3, 32, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv2d(32, 64, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+        )
+        flat = 128
+        self.proj = nn.Linear(flat, out_dim)
+
+    def forward(self, views: torch.Tensor) -> torch.Tensor:
+        # views: [B, V, 3, H, W] — mean pool over views
+        b, v, c, h, w = views.shape
+        x = views.reshape(b * v, c, h, w)
+        feats = self.backbone(x).reshape(b, v, -1).mean(dim=1)
+        return self.proj(feats)
+
+
+class MeshAIHybridTrainBundle(nn.Module):
+    """Geometry voxel + texture views + TRELLIS→Hunyuan latent bridge."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.geometry = GeometryVoxelHead()
+        self.texture = TextureViewEncoder()
+        self.bridge = TrellisHunyuanLatentAdapter(
+            trellis_dim=GEOM_IN_DIM,
+            hunyuan_dim=TEXTURE_LATENT_DIM,
+            hidden_dim=512,
+            depth=3,
+        )
+
+    def forward(self, geom_in: torch.Tensor, views: torch.Tensor) -> dict[str, torch.Tensor]:
+        voxel_pred = self.geometry(geom_in)
+        tex_latent = self.texture(views)
+        bridge_out = self.bridge(geom_in)
+        return {
+            "voxel_pred": voxel_pred,
+            "tex_latent": tex_latent,
+            "bridge_out": bridge_out,
+        }