Faz2: code/meshai_train/faz2_models.py

code/meshai_train/faz2_models.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""Faz 2: hibrit kopru + gercek TRELLIS/Hunyuan safetensors uzerinde LoRA."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from meshai_train.base_weights import (
+    BaseLoRATower,
+    _load_safetensors,
+    pick_lora_targets,
+)
+from meshai_train.models import GEOM_IN_DIM, TEXTURE_LATENT_DIM, MeshAIHybridTrainBundle
+
+FAZ2_VERSION = "v5.0-faz2-lora-base"
+
+
+class MeshAIFaz2Bundle(nn.Module):
+    """
+    Faz1 hibrit (geometry/texture/bridge) + TRELLIS/Hunyuan LoRA kuleleri.
+    Base agirliklar frozen; sadece LoRA + proj + hibrit egitilir.
+    """
+
+    def __init__(
+        self,
+        *,
+        trellis_targets: list[tuple[str, torch.Tensor]],
+        hunyuan_targets: list[tuple[str, torch.Tensor]],
+        lora_rank: int = 8,
+    ) -> None:
+        super().__init__()
+        self.hybrid = MeshAIHybridTrainBundle()
+        self.trellis_tower = BaseLoRATower(
+            trellis_targets,
+            input_dim=GEOM_IN_DIM,
+            rank=lora_rank,
+            out_dim=TEXTURE_LATENT_DIM,
+        )
+        # Hunyuan: view embedding -> tower
+        self.view_pool = nn.Sequential(
+            nn.Conv2d(3, 32, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+        )
+        self.view_proj = nn.Linear(32, GEOM_IN_DIM)
+        self.hunyuan_tower = BaseLoRATower(
+            hunyuan_targets,
+            input_dim=GEOM_IN_DIM,
+            rank=lora_rank,
+            out_dim=TEXTURE_LATENT_DIM,
+        )
+
+    def forward(self, geom_in: torch.Tensor, views: torch.Tensor) -> dict[str, torch.Tensor]:
+        hy = self.hybrid(geom_in, views)
+        trellis_feat = self.trellis_tower(geom_in)
+
+        b, v, c, h, w = views.shape
+        pooled = self.view_pool(views.reshape(b * v, c, h, w)).reshape(b, v, -1).mean(dim=1)
+        view_cond = self.view_proj(pooled)
+        hunyuan_feat = self.hunyuan_tower(view_cond)
+
+        return {
+            **hy,
+            "trellis_feat": trellis_feat,
+            "hunyuan_feat": hunyuan_feat,
+        }
+
+
+def faz2_loss(
+    out: dict[str, torch.Tensor],
+    batch: dict[str, torch.Tensor],
+) -> tuple[torch.Tensor, dict[str, float]]:
+    voxel_loss = nn.functional.mse_loss(out["voxel_pred"], batch["voxel_tgt"])
+    bridge_loss = nn.functional.mse_loss(out["bridge_out"], out["tex_latent"].detach())
+    # Gercek base kosullama: LoRA ciktilari hibrit latente hizalansin
+    trellis_align = nn.functional.mse_loss(out["trellis_feat"], out["bridge_out"].detach())
+    hunyuan_align = nn.functional.mse_loss(out["hunyuan_feat"], out["tex_latent"].detach())
+    tex_reg = out["tex_latent"].pow(2).mean() * 1e-4
+    total = voxel_loss + bridge_loss + 0.5 * trellis_align + 0.5 * hunyuan_align + tex_reg
+    return total, {
+        "voxel": float(voxel_loss.item()),
+        "bridge": float(bridge_loss.item()),
+        "trellis_align": float(trellis_align.item()),
+        "hunyuan_align": float(hunyuan_align.item()),
+        "tex_reg": float(tex_reg.item()),
+    }
+
+
+def build_faz2_from_weight_files(
+    weight_paths: dict[str, Path],
+    *,
+    lora_rank: int = 8,
+    trellis_mats: int = 4,
+    hunyuan_mats: int = 4,
+    log_fn: Any = print,
+) -> MeshAIFaz2Bundle:
+    trellis_targets: list[tuple[str, torch.Tensor]] = []
+    hunyuan_targets: list[tuple[str, torch.Tensor]] = []
+
+    for rel, path in weight_paths.items():
+        state = _load_safetensors(path)
+        picks = pick_lora_targets(state, max_matrices=trellis_mats if "TRELLIS" in rel else hunyuan_mats)
+        log_fn(f"[faz2] {rel}: {len(state)} tensor, LoRA aday={len(picks)}")
+        if "TRELLIS" in rel:
+            trellis_targets.extend(picks)
+        else:
+            hunyuan_targets.extend(picks)
+
+    if not trellis_targets:
+        raise RuntimeError("TRELLIS LoRA hedefi bulunamadi")
+    if not hunyuan_targets:
+        raise RuntimeError("Hunyuan LoRA hedefi bulunamadi")
+
+    # En buyuk N
+    trellis_targets = sorted(trellis_targets, key=lambda x: x[1].numel(), reverse=True)[:trellis_mats]
+    hunyuan_targets = sorted(hunyuan_targets, key=lambda x: x[1].numel(), reverse=True)[:hunyuan_mats]
+    log_fn(
+        f"[faz2] secilen TRELLIS mats={len(trellis_targets)} "
+        f"Hunyuan mats={len(hunyuan_targets)} rank={lora_rank}"
+    )
+    return MeshAIFaz2Bundle(
+        trellis_targets=trellis_targets,
+        hunyuan_targets=hunyuan_targets,
+        lora_rank=lora_rank,
+    )
+
+
+def save_faz2_checkpoint(
+    path: Path,
+    *,
+    epoch: int,
+    global_step: int,
+    model: MeshAIFaz2Bundle,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    payload = {
+        "version": FAZ2_VERSION,
+        "epoch": epoch,
+        "global_step": global_step,
+        "model": model.state_dict(),
+        "extra": extra or {},
+    }
+    torch.save(payload, path)
+
+
+def load_faz2_checkpoint(path: Path, model: MeshAIFaz2Bundle, device: str) -> int:
+    state = torch.load(path, map_location=device, weights_only=False)
+    if state.get("version") != FAZ2_VERSION:
+        # Faz1 hibrit ckpt ise sadece hybrid yukle
+        if "geometry" in state and "texture" in state and "bridge" in state:
+            model.hybrid.geometry.load_state_dict(state["geometry"])
+            model.hybrid.texture.load_state_dict(state["texture"])
+            model.hybrid.bridge.load_state_dict(state["bridge"])
+            return int(state.get("global_step", 0))
+        return 0
+    model.load_state_dict(state["model"], strict=False)
+    return int(state.get("global_step", 0))