code/meshai_train/engine.py

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.")