File size: 2,494 Bytes
1b93e89 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | 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,
}
|