#!/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}")