File size: 5,706 Bytes
687b68c | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | #!/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))
|