Faz2: code/meshai_train/faz2_engine.py

code/meshai_train/faz2_engine.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Faz 2 egitim motoru: base safetensors LoRA + preprocessed veri."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Subset
+
+from meshai_train.base_weights import ensure_base_weights
+from meshai_train.dataset import PreprocessedMeshDataset, collate_preprocessed
+from meshai_train.faz2_models import (
+    FAZ2_VERSION,
+    build_faz2_from_weight_files,
+    faz2_loss,
+    load_faz2_checkpoint,
+    save_faz2_checkpoint,
+)
+
+
+def start_faz2_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 = 50,
+    lora_rank: int = 8,
+    base_cache: Path | None = None,
+) -> None:
+    if not token:
+        raise RuntimeError("Faz2 icin HF_TOKEN gerekli (base weight + data)")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    cache = base_cache or Path("data/base_models")
+    log_fn(f"Pipeline surumu: {FAZ2_VERSION}")
+    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")
+
+    weight_paths = ensure_base_weights(token=token, cache_dir=cache, log_fn=log_fn)
+    model = build_faz2_from_weight_files(weight_paths, lora_rank=lora_rank, log_fn=log_fn)
+    model = model.to(device)
+
+    # Faz1 hibrit ckpt varsa hybrid'e yukle
+    faz1 = checkpoint_dir / "latest_model.pt"
+    if resume_from and resume_from.exists():
+        pass
+    elif faz1.exists():
+        resume_from = faz1
+
+    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=1 if device == "cuda" else 1,
+        shuffle=True,
+        pin_memory=device == "cuda",
+        collate_fn=collate_preprocessed,
+    )
+    val_loader = DataLoader(
+        Subset(dataset, val_idx),
+        batch_size=1,
+        shuffle=False,
+        pin_memory=device == "cuda",
+        collate_fn=collate_preprocessed,
+    )
+    log_fn(f"Veri: {len(train_idx)} train + {len(val_idx)} val | LoRA rank={lora_rank}")
+
+    # Egitilebilir: LoRA + proj + hybrid (frozen W buffer)
+    trainable = [p for p in model.parameters() if p.requires_grad]
+    opt = torch.optim.AdamW(trainable, lr=5e-5, weight_decay=0.01)
+    latest = checkpoint_dir / "latest_faz2_model.pt"
+    global_step = 0
+    if resume_from and resume_from.exists():
+        global_step = load_faz2_checkpoint(resume_from, model, device)
+        log_fn(f"Resume: {resume_from} step={global_step}")
+
+    def _eval() -> 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, _ = faz2_loss(out, batch)
+                losses.append(float(loss.item()))
+        model.train()
+        return float(np.mean(losses)) if losses else float("nan")
+
+    trainable_n = sum(p.numel() for p in trainable)
+    frozen_n = sum(p.numel() for p in model.parameters() if not p.requires_grad)
+    buffer_n = sum(b.numel() for b in model.buffers())
+    log_fn(f"Param: trainable={trainable_n:,} frozen={frozen_n:,} buffers(W)={buffer_n:,}")
+
+    for epoch in range(1, epochs + 1):
+        log_fn(f"--- Faz2 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 = faz2_loss(out, batch)
+            if not torch.isfinite(loss):
+                monitor.note_nan_skip("faz2")
+                continue
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(trainable, 1.0)
+            opt.step()
+            global_step += 1
+            epoch_losses.append(float(loss.item()))
+            monitor.note_step(global_step, "faz2", parts["voxel"])
+            if global_step <= 3 or global_step % 20 == 0:
+                log_fn(
+                    f"step={global_step} loss={float(loss.item()):.6f} "
+                    f"voxel={parts['voxel']:.4f} trellis={parts['trellis_align']:.4f} "
+                    f"hunyuan={parts['hunyuan_align']:.4f}"
+                )
+            if checkpoint_every > 0 and global_step % checkpoint_every == 0:
+                save_faz2_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
+                # Ayrica latest_model.pt olarak da yaz (orchestrator uyumu)
+                save_faz2_checkpoint(
+                    checkpoint_dir / "latest_model.pt",
+                    epoch=epoch,
+                    global_step=global_step,
+                    model=model,
+                )
+                log_fn(
+                    f"CHECKPOINT_SAVED step={global_step} -> {latest} "
+                    f"({latest.stat().st_size // (1024 * 1024)} MB)"
+                )
+            if global_step % validation_every == 0:
+                val = _eval()
+                monitor.note_validation(global_step, val, val)
+                log_vram_fn(f"step_{global_step}")
+
+        val = _eval()
+        mean = float(np.mean(epoch_losses)) if epoch_losses else float("nan")
+        monitor.note_epoch_end(epoch, epochs, mean, val, val, val)
+        save_faz2_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
+        save_faz2_checkpoint(
+            checkpoint_dir / "latest_model.pt",
+            epoch=epoch,
+            global_step=global_step,
+            model=model,
+        )
+        clear_gpu_fn()
+        log_fn(f"Epoch {epoch} kaydedildi -> {latest} ({latest.stat().st_size // (1024 * 1024)} MB)")
+        log_fn(f"CHECKPOINT_SAVED step={global_step} epoch={epoch} -> {latest}")
+
+    monitor.finish(ok=True)
+    log_fn("Faz2 TRELLIS/Hunyuan LoRA egitimi tamamlandi.")
+    log_fn(f"Cikti: {latest}")