from __future__ import annotations import json from pathlib import Path from typing import Any import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, Subset from meshai_train.dataset import PreprocessedMeshDataset, collate_preprocessed from meshai_train.models import MeshAIHybridTrainBundle TRAIN_VERSION = "v4.0-real-preprocessed" def hybrid_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()) tex_reg = out["tex_latent"].pow(2).mean() * 1e-4 total = voxel_loss + bridge_loss + tex_reg return total, { "voxel": float(voxel_loss.item()), "bridge": float(bridge_loss.item()), "tex_reg": float(tex_reg.item()), } def save_real_checkpoint( path: Path, *, epoch: int, global_step: int, model: MeshAIHybridTrainBundle, extra: dict[str, Any] | None = None, ) -> None: payload = { "version": TRAIN_VERSION, "epoch": epoch, "global_step": global_step, "geometry": model.geometry.state_dict(), "texture": model.texture.state_dict(), "bridge": model.bridge.state_dict(), "extra": extra or {}, } torch.save(payload, path) def load_real_checkpoint(path: Path, model: MeshAIHybridTrainBundle, device: str) -> int: state = torch.load(path, map_location=device, weights_only=False) if state.get("version") != TRAIN_VERSION: return 0 model.geometry.load_state_dict(state["geometry"]) model.texture.load_state_dict(state["texture"]) model.bridge.load_state_dict(state["bridge"]) return int(state.get("global_step", 0)) def start_real_training( *, monitor: Any, checkpoint_dir: Path, output_dir: Path, token: str | None, epochs: int, resume_from: Path | None, validation_every: int, val_ratio: float, limit: int | None, hf_repo: str, data_root: Path | None, log_fn: Any, log_vram_fn: Any, clear_gpu_fn: Any, load_val_uids_fn: Any, checkpoint_every: int = 100, ) -> None: device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float32 log_fn(f"Pipeline surumu: {TRAIN_VERSION} (gercek preprocessed latent + render)") log_fn(f"Veri: {hf_repo}" + (f" | limit={limit}" if limit else "")) if torch.cuda.is_available(): log_fn(f"GPU: {torch.cuda.get_device_name(0)}") log_vram_fn("startup") dataset = PreprocessedMeshDataset( token=token, data_root=data_root, hf_repo=hf_repo, limit=limit, ) val_uids = load_val_uids_fn() val_idx = [i for i, o in enumerate(dataset.objects) if str(o.get("uid")) in val_uids] if not val_idx: val_count = max(1, int(len(dataset) * val_ratio)) val_idx = list(range(val_count)) train_idx = [i for i in range(len(dataset)) if i not in set(val_idx)] if not train_idx: train_idx = list(range(len(dataset))) val_idx = train_idx[:1] train_loader = DataLoader( Subset(dataset, train_idx), batch_size=2 if device == "cuda" else 1, shuffle=True, pin_memory=device == "cuda", collate_fn=collate_preprocessed, ) val_loader = DataLoader( Subset(dataset, val_idx), batch_size=2 if device == "cuda" else 1, shuffle=False, pin_memory=device == "cuda", collate_fn=collate_preprocessed, ) log_fn(f"Veri: {len(train_idx)} train + {len(val_idx)} val (gercek latent/render).") model = MeshAIHybridTrainBundle().to(device=device, dtype=dtype) opt = torch.optim.AdamW(model.parameters(), lr=1e-4, fused=False) latest = checkpoint_dir / "latest_model.pt" global_step = 0 if resume_from and resume_from.exists(): global_step = load_real_checkpoint(resume_from, model, device) log_fn(f"Resume: {resume_from} step={global_step}") def _eval() -> tuple[float, float]: model.eval() losses: list[float] = [] with torch.no_grad(): for batch in val_loader: batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} out = model(batch["geom_in"], batch["views"]) loss, _ = hybrid_loss(out, batch) losses.append(float(loss.item())) model.train() mean = float(np.mean(losses)) if losses else float("nan") return mean, mean for epoch in range(1, epochs + 1): log_fn(f"--- Epoch {epoch}/{epochs} ---") model.train() epoch_losses: list[float] = [] for batch in train_loader: batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} opt.zero_grad(set_to_none=True) out = model(batch["geom_in"], batch["views"]) loss, parts = hybrid_loss(out, batch) if not torch.isfinite(loss): monitor.note_nan_skip("hybrid") continue loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() global_step += 1 epoch_losses.append(float(loss.item())) monitor.note_step(global_step, "geometry", parts["voxel"]) monitor.last["train_loss_texture"] = round(parts["bridge"], 6) if global_step <= 3 or global_step % 25 == 0: log_fn( f"step={global_step} loss={float(loss.item()):.6f} " f"voxel={parts['voxel']:.6f} bridge={parts['bridge']:.6f}" ) if checkpoint_every > 0 and global_step % checkpoint_every == 0: save_real_checkpoint(latest, epoch=epoch, global_step=global_step, model=model) log_fn( f"CHECKPOINT_SAVED step={global_step} -> {latest} " f"({latest.stat().st_size // 1024} KB)" ) if global_step % validation_every == 0: val_geom, val_tex = _eval() monitor.note_validation(global_step, val_geom, val_tex) log_vram_fn(f"step_{global_step}") val_geom, val_tex = _eval() geom_mean = float(np.mean(epoch_losses)) if epoch_losses else float("nan") monitor.note_epoch_end(epoch, epochs, geom_mean, val_tex, val_geom, val_tex) save_real_checkpoint(latest, epoch=epoch, global_step=global_step, model=model) clear_gpu_fn() log_fn(f"Epoch {epoch} kaydedildi -> {latest} ({latest.stat().st_size // 1024} KB)") log_fn(f"CHECKPOINT_SAVED step={global_step} epoch={epoch} -> {latest}") monitor.finish(ok=True) log_fn("Gercek preprocessed egitim tamamlandi.") log_fn("Not: Tam TRELLIS/Hunyuan agirlik fine-tune sonraki adim; bu asama MeshAI hibrit latent koprusu.")