File size: 4,388 Bytes
28eb85c | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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"],
}
|