| 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:
|
|
|
| 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,
|
| }
|
|
|
