Faz2 dep: code/meshai_bridge/latent_adapter.py

code/meshai_bridge/latent_adapter.py

@@ -0,0 +1,126 @@
+"""Latent adapter R&D scaffold (requires real TRELLIS/Hunyuan model tensors)."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+try:
+    import torch
+    import torch.nn as nn
+except ImportError:  # pragma: no cover - optional for preprocess-only runs
+    torch = None  # type: ignore
+    nn = None  # type: ignore
+
+
+def _utc_now() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+@dataclass
+class LatentShapeLogger:
+    """Log tensor shapes until real model hooks are wired."""
+
+    log_path: Path | None = None
+    records: list[dict[str, Any]] = field(default_factory=list)
+
+    def log(self, name: str, tensor: Any, *, source: str = "unknown") -> dict[str, Any]:
+        if torch is not None and isinstance(tensor, torch.Tensor):
+            shape = list(tensor.shape)
+            dtype = str(tensor.dtype)
+            device = str(tensor.device)
+        elif hasattr(tensor, "shape"):
+            shape = list(getattr(tensor, "shape"))
+            dtype = str(getattr(tensor, "dtype", "?"))
+            device = "numpy"
+        else:
+            shape = []
+            dtype = type(tensor).__name__
+            device = "?"
+
+        entry = {
+            "name": name,
+            "source": source,
+            "shape": shape,
+            "dtype": dtype,
+            "device": device,
+            "logged_at": _utc_now(),
+        }
+        self.records.append(entry)
+        if self.log_path is not None:
+            self.log_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(self.log_path, "a", encoding="utf-8") as handle:
+                handle.write(json.dumps(entry, ensure_ascii=False) + "\n")
+        return entry
+
+    def summary(self) -> dict[str, Any]:
+        return {
+            "count": len(self.records),
+            "shapes": {r["name"]: r["shape"] for r in self.records},
+            "updated_at": _utc_now(),
+        }
+
+
+if nn is not None:
+
+    class TrellisHunyuanLatentAdapter(nn.Module):
+        """Small MLP adapter: TRELLIS structured latent -> Hunyuan conditioning latent."""
+
+        def __init__(
+            self,
+            trellis_dim: int,
+            hunyuan_dim: int,
+            *,
+            hidden_dim: int = 512,
+            depth: int = 2,
+            adapter_type: str = "mlp",
+        ) -> None:
+            super().__init__()
+            self.adapter_type = adapter_type
+            self.trellis_dim = trellis_dim
+            self.hunyuan_dim = hunyuan_dim
+            layers: list[nn.Module] = []
+            in_dim = trellis_dim
+            for _ in range(max(depth - 1, 0)):
+                layers.extend([nn.Linear(in_dim, hidden_dim), nn.GELU()])
+                in_dim = hidden_dim
+            layers.append(nn.Linear(in_dim, hunyuan_dim))
+            self.net = nn.Sequential(*layers)
+            self.logger = LatentShapeLogger()
+
+        def forward(self, trellis_latent: "torch.Tensor") -> "torch.Tensor":
+            self.logger.log("trellis_latent_in", trellis_latent, source="trellis")
+            flat = trellis_latent.reshape(trellis_latent.shape[0], -1)
+            if flat.shape[-1] != self.trellis_dim:
+                raise ValueError(
+                    f"Expected trellis_dim={self.trellis_dim}, got {flat.shape[-1]}"
+                )
+            out = self.net(flat)
+            self.logger.log("hunyuan_latent_out", out, source="adapter")
+            return out
+
+else:
+
+    class TrellisHunyuanLatentAdapter:  # type: ignore[no-redef]
+        def __init__(self, *args, **kwargs) -> None:
+            raise ImportError("torch required for TrellisHunyuanLatentAdapter")
+
+
+def research_notes() -> dict[str, Any]:
+    """Document planned losses and integration points for future real-model wiring."""
+    return {
+        "status": "rd_only",
+        "requires": [
+            "trellis_structured_latent_export",
+            "hunyuan_denoising_conditioning_hook",
+        ],
+        "candidate_losses": [
+            "texture_view_reconstruction",
+            "uv_consistency",
+            "normal_curvature_texture_consistency",
+        ],
+        "adapter_variants": ["linear", "mlp", "cross_attention"],
+    }