| """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:
|
| torch = None
|
| nn = None
|
|
|
|
|
| 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:
|
| 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"],
|
| }
|
|
|
