|
|
| """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 = 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}")
|
|
|
|
|
| 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)
|
|
|
| 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}")
|
|
|
