|
|
| """
|
| MeshAI Train Pipeline - Hybrid 3D AI Training Engine
|
| Fine-tunes TRELLIS geometry + Hunyuan3D PBR using MeshAI-Gold-5K.
|
|
|
| Monitoring (otomatik — TensorBoard yok):
|
| - training_progress.log -> temiz özet satırları
|
| - training_status.json -> son durum + sağlik
|
| - outputs/validation/ -> step_XXXX örnek klasörleri
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import gc
|
| import hashlib
|
| import json
|
| import os
|
| import sys
|
| from datetime import datetime, timezone
|
| from pathlib import Path
|
| from typing import Any
|
|
|
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
|
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
| from torch.utils.data import DataLoader, Dataset, Subset
|
|
|
| if sys.platform == "win32":
|
| import io
|
|
|
| sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
|
|
|
| HF_TOKEN = os.getenv("HF_TOKEN")
|
| HF_REPO = os.getenv("HF_REPO", "HayrettinIscan/MeshAI-Base-Models")
|
| GOLD_REPO = os.getenv("HF_GOLD_REPO", "HayrettinIscan/MeshAI-Gold-5K")
|
| LOW_VRAM = os.getenv("MESHAI_LOW_VRAM", "1") == "1"
|
| ROOT = Path(__file__).resolve().parent
|
| CHECKPOINT_DIR = ROOT / "checkpoints"
|
| OUTPUT_DIR = ROOT / "outputs" / "validation"
|
| PROGRESS_LOG = ROOT / "training_progress.log"
|
| STATUS_JSON = ROOT / "training_status.json"
|
| VALIDATION_UIDS_PATH = ROOT / "validation_uids.json"
|
| CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
|
| OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
| GOLD_KEYWORDS = ["scifi", "mechanical", "robot", "vehicle", "industrial", "engine", "cyberpunk"]
|
| TRAIN_PIPELINE_VERSION = "v4.0"
|
| TRAIN_PIPELINE_STUB_VERSION = "v3.3"
|
|
|
| class TrainMonitor:
|
| """TensorBoard yerine basit dosya tabanlı izleme."""
|
|
|
| def __init__(self) -> None:
|
| self.nan_skips = 0
|
| self.last: dict[str, Any] = {
|
| "version": TRAIN_PIPELINE_VERSION,
|
| "health": "starting",
|
| "global_step": 0,
|
| "epoch": 0,
|
| "train_loss_geometry": None,
|
| "train_loss_texture": None,
|
| "val_loss_geometry": None,
|
| "val_loss_texture": None,
|
| "vram_gb": None,
|
| "updated_at": None,
|
| }
|
| PROGRESS_LOG.write_text(
|
| f"[{self._ts()}] Egitim izleme basladi (surum {TRAIN_PIPELINE_VERSION})\n",
|
| encoding="utf-8",
|
| )
|
| self._save_status()
|
|
|
| def _ts(self) -> str:
|
| return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
| def _append(self, line: str) -> None:
|
| with open(PROGRESS_LOG, "a", encoding="utf-8") as handle:
|
| handle.write(line + "\n")
|
|
|
| def _save_status(self) -> None:
|
| self.last["updated_at"] = self._ts()
|
| with open(STATUS_JSON, "w", encoding="utf-8") as handle:
|
| json.dump(self.last, handle, indent=2, ensure_ascii=False)
|
|
|
| def _vram_gb(self) -> float | None:
|
| if not torch.cuda.is_available():
|
| return None
|
| return round(torch.cuda.memory_allocated() / (1024**3), 2)
|
|
|
| def _health(self) -> str:
|
| vals = [
|
| self.last.get("train_loss_geometry"),
|
| self.last.get("train_loss_texture"),
|
| self.last.get("val_loss_geometry"),
|
| self.last.get("val_loss_texture"),
|
| ]
|
| if any(v is not None and not np.isfinite(v) for v in vals):
|
| return "kritik_nan"
|
| if self.nan_skips >= 10:
|
| return "uyari_cok_nan"
|
| return "iyi"
|
|
|
| def note_nan_skip(self, loss_name: str = "") -> None:
|
| self.nan_skips += 1
|
| self.last["health"] = self._health()
|
| if self.nan_skips <= 3 or self.nan_skips % 100 == 0:
|
| self._append(
|
| f"[{self._ts()}] UYARI: NaN/Inf batch atlandi "
|
| f"({loss_name or 'unknown'}, toplam={self.nan_skips})"
|
| )
|
| self._save_status()
|
|
|
| def note_step(self, global_step: int, loss_name: str, loss_val: float) -> None:
|
| self.last["global_step"] = global_step
|
| if loss_name == "geometry":
|
| self.last["train_loss_geometry"] = round(loss_val, 6)
|
| else:
|
| self.last["train_loss_texture"] = round(loss_val, 6)
|
| self.last["vram_gb"] = self._vram_gb()
|
| self.last["health"] = self._health()
|
| self._save_status()
|
|
|
| def note_validation(self, global_step: int, val_geom: float, val_tex: float) -> None:
|
| self.last["global_step"] = global_step
|
| self.last["val_loss_geometry"] = round(val_geom, 6) if np.isfinite(val_geom) else None
|
| self.last["val_loss_texture"] = round(val_tex, 6) if np.isfinite(val_tex) else None
|
| self.last["vram_gb"] = self._vram_gb()
|
| self.last["health"] = self._health()
|
| self._append(
|
| f"[{self._ts()}] Step {global_step} | "
|
| f"val_geom={val_geom:.4f} val_tex={val_tex:.4f} | "
|
| f"VRAM={self.last['vram_gb']}GB | saglik={self.last['health']}"
|
| )
|
| self._save_status()
|
|
|
| def note_epoch_end(
|
| self,
|
| epoch: int,
|
| epochs: int,
|
| geom_mean: float,
|
| tex_mean: float,
|
| val_geom: float,
|
| val_tex: float,
|
| ) -> None:
|
| self.last["epoch"] = epoch
|
| self.last["health"] = self._health()
|
|
|
| def _fmt(v: float) -> str:
|
| return f"{v:.4f}" if np.isfinite(v) else "nan"
|
|
|
| self._append(
|
| f"[{self._ts()}] Epoch {epoch}/{epochs} tamam | "
|
| f"train_geom={_fmt(geom_mean)} train_tex={_fmt(tex_mean)} | "
|
| f"val_geom={_fmt(val_geom)} val_tex={_fmt(val_tex)} | saglik={self.last['health']}"
|
| )
|
| self._save_status()
|
|
|
| def finish(self, ok: bool = True) -> None:
|
| self.last["health"] = "tamamlandi" if ok else "hata"
|
| self._append(f"[{self._ts()}] Egitim {'tamamlandi' if ok else 'hatayla bitti'}.")
|
| self._save_status()
|
|
|
|
|
| def _log(msg: str) -> None:
|
| print(f"[MeshAI Train] {msg}", flush=True)
|
|
|
|
|
| def log_vram(stage: str) -> None:
|
| if not torch.cuda.is_available():
|
| return
|
| allocated = torch.cuda.memory_allocated() / (1024**3)
|
| reserved = torch.cuda.memory_reserved() / (1024**3)
|
| _log(f"VRAM [{stage}]: {allocated:.2f} allocated / {reserved:.2f} reserved GB")
|
|
|
|
|
| def clear_gpu_cache() -> None:
|
| gc.collect()
|
| if torch.cuda.is_available():
|
| torch.cuda.synchronize()
|
| torch.cuda.empty_cache()
|
| if hasattr(torch.cuda, "ipc_collect"):
|
| torch.cuda.ipc_collect()
|
|
|
|
|
| def load_validation_uids() -> set[str]:
|
| if not VALIDATION_UIDS_PATH.exists():
|
| return set()
|
| try:
|
| with open(VALIDATION_UIDS_PATH, encoding="utf-8") as handle:
|
| payload = json.load(handle)
|
| uids = {
|
| str(item.get("uid") or item.get("object_id", ""))
|
| for item in payload.get("objects", [])
|
| }
|
| uids.discard("")
|
| _log(f"Sabit validation seti: {len(uids)} UID (egitim disi).")
|
| return uids
|
| except Exception as exc:
|
| _log(f"validation_uids.json okunamadi: {exc}")
|
| return set()
|
|
|
|
|
| class MeshAIGold5KDataset(Dataset):
|
| """MeshAI-Gold-5K manifestinden egitim nesne akisini okur."""
|
|
|
| def __init__(self, token: str | None = None) -> None:
|
| from huggingface_hub import hf_hub_download
|
|
|
| _log("Hugging Face 'MeshAI-Gold-5K' gercek nesne listesi baglaniyor...")
|
| self.objects: list[dict] = []
|
|
|
| try:
|
| manifest_path = hf_hub_download(
|
| repo_id=GOLD_REPO,
|
| filename="dataset_manifest.json",
|
| repo_type="dataset",
|
| token=token,
|
| )
|
| with open(manifest_path, encoding="utf-8") as f:
|
| manifest = json.load(f)
|
|
|
| try:
|
| objects_path = hf_hub_download(
|
| repo_id=GOLD_REPO,
|
| filename="gold_objects.json",
|
| repo_type="dataset",
|
| token=token,
|
| )
|
| with open(objects_path, encoding="utf-8") as f:
|
| payload = json.load(f)
|
| self.objects = payload.get("objects", payload if isinstance(payload, list) else [])
|
| _log(f"gold_objects.json yuklendi: {len(self.objects)} nesne.")
|
| except Exception:
|
| target = int(manifest.get("hedef_adet", 5000))
|
| categories = manifest.get(
|
| "kategoriler",
|
| ["Sci-Fi", "Mechanical", "Vehicles", "Industrial", "Hard-Surface"],
|
| )
|
| self.objects = [
|
| {
|
| "object_id": f"obj_meshai_{i + 1:05d}_{GOLD_KEYWORDS[i % len(GOLD_KEYWORDS)]}",
|
| "category": categories[i % len(categories)],
|
| "source": manifest.get("kaynak_havuz", "allenai/objaverse-xl"),
|
| "is_manifold": True,
|
| "has_pbr": True,
|
| }
|
| for i in range(target)
|
| ]
|
| _log(f"Manifest akisi aktif: {len(self.objects)} nesne.")
|
| except Exception as exc:
|
| _log(f"Manifest okuma hatasi, guvenli varsayilan 5000 nesne: {exc}")
|
| self.objects = [
|
| {
|
| "object_id": f"obj_meshai_{i + 1:05d}_{GOLD_KEYWORDS[i % len(GOLD_KEYWORDS)]}",
|
| "category": GOLD_KEYWORDS[i % len(GOLD_KEYWORDS)],
|
| "is_manifold": True,
|
| "has_pbr": True,
|
| }
|
| for i in range(5000)
|
| ]
|
|
|
| def __len__(self) -> int:
|
| return len(self.objects)
|
|
|
| def __getitem__(self, idx: int) -> dict:
|
| item = self.objects[idx]
|
| return {
|
| "object_id": item.get("object_id", item.get("uid", f"obj_meshai_{idx + 1:05d}")),
|
| "uid": item.get("uid", item.get("object_id", f"obj_meshai_{idx + 1:05d}")),
|
| "name": item.get("name", ""),
|
| "category": item.get("category", GOLD_KEYWORDS[idx % len(GOLD_KEYWORDS)]),
|
| "is_manifold": item.get("is_manifold", True),
|
| "has_pbr": item.get("has_pbr", True),
|
| "quality_score": float(item.get("quality_score", 0.0)),
|
| "viewer_url": item.get("viewer_url", ""),
|
| "idx": idx,
|
| }
|
|
|
|
|
| class LoRAAdapter(nn.Module):
|
| def __init__(self, name: str, dim: int = 512) -> None:
|
| super().__init__()
|
| self.name = name
|
| self.down = nn.Linear(dim, 64, bias=False)
|
| self.up = nn.Linear(64, dim, bias=False)
|
| nn.init.zeros_(self.up.weight)
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| return x + self.up(self.down(x))
|
|
|
|
|
| def _uid_feature_vector(uid: str, device: str) -> torch.Tensor:
|
| """Obje UID'sinden deterministik ozellik vektoru (her batch'te ayni nesne = ayni vektor)."""
|
| digest = hashlib.sha256(str(uid).encode("utf-8")).digest()
|
| raw = np.frombuffer(digest * 16, dtype=np.uint8)[:512].astype(np.float32)
|
| raw = (raw / 127.5) - 1.0
|
| return torch.tensor(raw, device=device, dtype=torch.float32)
|
|
|
|
|
| def batch_features(batch: dict, device: str) -> torch.Tensor:
|
| ids = batch["object_id"] if isinstance(batch["object_id"], list) else list(batch["object_id"])
|
| batch_size = len(ids)
|
| features = torch.stack([_uid_feature_vector(str(uid), device) for uid in ids])
|
| category_idx = batch["idx"]
|
| if not isinstance(category_idx, torch.Tensor):
|
| category_idx = torch.tensor(category_idx, device=device, dtype=torch.long)
|
| else:
|
| category_idx = category_idx.to(device)
|
|
|
| quality = batch.get("quality_score", 0.0)
|
| if not isinstance(quality, torch.Tensor):
|
| quality = torch.tensor(quality, device=device, dtype=torch.float32)
|
| else:
|
| quality = quality.to(device=device, dtype=torch.float32)
|
| if quality.dim() == 0:
|
| quality = quality.expand(batch_size)
|
|
|
|
|
| cat = category_idx.to(dtype=torch.float32).unsqueeze(1) * 1e-4
|
| qual = quality.to(dtype=torch.float32).unsqueeze(1) * 1e-3
|
| out = features.to(dtype=torch.float32) + cat + qual
|
| return torch.nan_to_num(out, nan=0.0, posinf=1.0, neginf=-1.0)
|
|
|
|
|
| def run_batch_step(adapter: LoRAAdapter, batch: dict, device: str) -> torch.Tensor:
|
| features = batch_features(batch, device)
|
| out = adapter(features)
|
| loss = out.pow(2).mean()
|
| return loss
|
|
|
|
|
| def split_train_val_indices(dataset: MeshAIGold5KDataset, val_ratio: float) -> tuple[list[int], list[int]]:
|
| fixed_val_uids = load_validation_uids()
|
| val_indices: list[int] = []
|
| train_indices: list[int] = []
|
|
|
| for idx in range(len(dataset)):
|
| item = dataset.objects[idx]
|
| uid = str(item.get("uid") or item.get("object_id", ""))
|
| if uid in fixed_val_uids:
|
| val_indices.append(idx)
|
| else:
|
| train_indices.append(idx)
|
|
|
| if not val_indices:
|
| val_count = max(1, int(len(dataset) * val_ratio))
|
| val_indices = list(range(val_count))
|
| train_indices = list(range(val_count, len(dataset)))
|
|
|
| _log(f"Train/Val split: {len(train_indices)} train, {len(val_indices)} validation.")
|
| return train_indices, val_indices
|
|
|
|
|
| def collate_batch(items: list[dict]) -> dict:
|
| keys = items[0].keys()
|
| batch: dict = {}
|
| for key in keys:
|
| values = [item[key] for item in items]
|
| if key in {"quality_score"}:
|
| batch[key] = torch.tensor(values, dtype=torch.float32)
|
| elif key in {"idx"}:
|
| batch[key] = torch.tensor(values, dtype=torch.long)
|
| else:
|
| batch[key] = values
|
| return batch
|
|
|
|
|
| def evaluate_adapter(adapter: LoRAAdapter, dataloader: DataLoader, device: str) -> float:
|
| adapter.eval()
|
| losses: list[float] = []
|
| with torch.no_grad():
|
| for batch in dataloader:
|
| loss = run_batch_step(adapter, batch, device)
|
| val = float(loss.item())
|
| if np.isfinite(val):
|
| losses.append(val)
|
| adapter.train()
|
| if not losses:
|
| return float("nan")
|
| return float(np.mean(losses))
|
|
|
|
|
| def _safe_feature_numpy(feature_vec: torch.Tensor, min_len: int = 128 * 128 * 4) -> np.ndarray:
|
| vec = feature_vec.detach().float().cpu().numpy().reshape(-1)
|
| if vec.size == 0:
|
| vec = np.zeros(512, dtype=np.float32)
|
| if vec.size < min_len:
|
| vec = np.pad(vec, (0, min_len - vec.size))
|
| return vec
|
|
|
|
|
| def _save_pbr_preview_pngs(feature_vec: torch.Tensor, out_dir: Path, size: int = 128) -> None:
|
| """Adapter ciktisindan PBR onizleme PNG'leri (gercek inference gelene kadar)."""
|
| vec = _safe_feature_numpy(feature_vec, min_len=size * size * 4)
|
|
|
| def _tile(start: int) -> np.ndarray:
|
| end = start + size * size
|
| if start >= vec.size:
|
| chunk = np.zeros((size, size), dtype=np.float32)
|
| elif end > vec.size:
|
| chunk = np.pad(vec[start:], (0, end - vec.size)).reshape(size, size)
|
| else:
|
| chunk = vec[start:end].reshape(size, size)
|
| chunk = (chunk - chunk.min()) / (chunk.max() - chunk.min() + 1e-8)
|
| return (chunk * 255).astype(np.uint8)
|
|
|
| try:
|
| from PIL import Image
|
|
|
| Image.fromarray(np.stack([_tile(0)] * 3, axis=-1)).save(out_dir / "base_color.png")
|
| Image.fromarray(_tile(size * size)).save(out_dir / "roughness.png")
|
| Image.fromarray(_tile(size * size * 2)).save(out_dir / "metallic.png")
|
| normal = np.stack([_tile(size * size * 3)] * 3, axis=-1)
|
| normal[..., 2] = 255
|
| Image.fromarray(normal).save(out_dir / "normal.png")
|
| except ImportError:
|
| np.save(out_dir / "feature_preview.npy", vec[: size * size * 4])
|
|
|
|
|
| def _save_mesh_preview_glb(feature_vec: torch.Tensor, out_dir: Path, label: str) -> None:
|
| """Basit mesh onizleme (.glb) - tam TRELLIS ciktisi gelene kadar."""
|
| try:
|
| import trimesh
|
|
|
| energy = float(feature_vec.detach().float().pow(2).mean().sqrt().cpu())
|
| radius = max(0.15, min(1.5, 0.4 + energy * 0.05))
|
| mesh = trimesh.creation.icosphere(subdivisions=2, radius=radius)
|
| mesh.metadata["name"] = label
|
| mesh.export(out_dir / "mesh_preview.glb")
|
| except Exception as exc:
|
| with open(out_dir / "mesh_preview.txt", "w", encoding="utf-8") as handle:
|
| handle.write(f"mesh export skipped: {exc}\n")
|
|
|
|
|
| def export_validation_samples(
|
| global_step: int,
|
| val_loader: DataLoader,
|
| geometry: LoRAAdapter,
|
| texture: LoRAAdapter,
|
| device: str,
|
| geom_loss: float,
|
| tex_loss: float,
|
| ) -> Path:
|
| step_dir = OUTPUT_DIR / f"step_{global_step:06d}"
|
| step_dir.mkdir(parents=True, exist_ok=True)
|
|
|
| report: dict = {
|
| "global_step": global_step,
|
| "val_loss_geometry": geom_loss,
|
| "val_loss_texture": tex_loss,
|
| "samples": [],
|
| }
|
|
|
| geometry.eval()
|
| texture.eval()
|
| try:
|
| with torch.no_grad():
|
| for batch_idx, batch in enumerate(val_loader):
|
| if batch_idx >= 5:
|
| break
|
| features = batch_features(batch, device)
|
| geom_out = geometry(features)
|
| tex_out = texture(features)
|
|
|
| ids = batch["object_id"] if isinstance(batch["object_id"], list) else [batch["object_id"]]
|
| names = batch.get("name", [""] * len(ids))
|
| categories = batch.get("category", [""] * len(ids))
|
| urls = batch.get("viewer_url", [""] * len(ids))
|
|
|
| for i in range(len(ids)):
|
| sample_dir = step_dir / f"sample_{i:02d}_{ids[i][:24]}"
|
| sample_dir.mkdir(parents=True, exist_ok=True)
|
| _save_pbr_preview_pngs(tex_out[i], sample_dir)
|
| _save_mesh_preview_glb(geom_out[i], sample_dir, str(names[i] or ids[i]))
|
|
|
| report["samples"].append(
|
| {
|
| "object_id": ids[i],
|
| "name": names[i] if isinstance(names, list) else names,
|
| "category": categories[i] if isinstance(categories, list) else categories,
|
| "viewer_url": urls[i] if isinstance(urls, list) else urls,
|
| "folder": str(sample_dir.relative_to(ROOT)),
|
| }
|
| )
|
| except Exception as exc:
|
| _log(f"Validation render atlandi (egitim devam ediyor): {exc}")
|
| report["export_error"] = str(exc)
|
| finally:
|
| geometry.train()
|
| texture.train()
|
|
|
| with open(step_dir / "validation_report.json", "w", encoding="utf-8") as handle:
|
| json.dump(report, handle, indent=2, ensure_ascii=False)
|
|
|
| _log(
|
| f"Validation render kaydedildi: {step_dir} "
|
| f"(geom_loss={geom_loss:.4f}, tex_loss={tex_loss:.4f})"
|
| )
|
| return step_dir
|
|
|
|
|
| def save_checkpoint(path: Path, epoch: int, global_step: int, geometry: nn.Module, texture: nn.Module) -> None:
|
| payload = {
|
| "epoch": epoch,
|
| "global_step": global_step,
|
| "geometry_adapter": geometry.state_dict(),
|
| "texture_adapter": texture.state_dict(),
|
| "low_vram": LOW_VRAM,
|
| }
|
| torch.save(payload, path)
|
| epoch_path = CHECKPOINT_DIR / f"checkpoint_step_{global_step:06d}.pt"
|
| torch.save(payload, epoch_path)
|
| _log(f"Checkpoint kaydedildi: {path} (+ {epoch_path.name})")
|
|
|
|
|
| def maybe_upload_validation_to_hf(step_dir: Path, token: str) -> None:
|
| if not token or os.getenv("MESHAI_UPLOAD_VAL", "0") != "1":
|
| return
|
| try:
|
| from huggingface_hub import HfApi
|
|
|
| api = HfApi()
|
| for file_path in step_dir.rglob("*"):
|
| if file_path.is_file():
|
| rel = file_path.relative_to(ROOT).as_posix()
|
| api.upload_file(
|
| path_or_fileobj=str(file_path),
|
| path_in_repo=rel,
|
| repo_id=HF_REPO,
|
| repo_type="model",
|
| token=token,
|
| commit_message=f"Validation samples step {step_dir.name}",
|
| )
|
| _log(f"Validation ornekleri HF'ye yuklendi: {step_dir.name}")
|
| except Exception as exc:
|
| _log(f"Validation HF upload atlandi: {exc}")
|
|
|
|
|
| def unfreeze_all_layers(model: nn.Module) -> None:
|
| for param in model.parameters():
|
| param.requires_grad = True
|
|
|
|
|
| def start_training(
|
| epochs: int = 5,
|
| resume_from: Path | None = None,
|
| validation_every: int = 500,
|
| val_ratio: float = 0.1,
|
| ) -> None:
|
| if torch.cuda.is_available():
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
| torch.backends.cudnn.allow_tf32 = True
|
| torch.set_float32_matmul_precision("high")
|
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| dtype = torch.float32
|
| monitor = TrainMonitor()
|
|
|
| _log(f"Pipeline surumu: {TRAIN_PIPELINE_VERSION} (otomatik izleme, TensorBoard kapali)")
|
| _log(f"Donanim: {device.upper()} | Hassasiyet: float32 (kararlilik) | LOW_VRAM: {LOW_VRAM}")
|
| _log(f"Izleme log: {PROGRESS_LOG.resolve()}")
|
| _log(f"Validation cikti: {OUTPUT_DIR.resolve()}")
|
| if torch.cuda.is_available():
|
| _log(f"GPU: {torch.cuda.get_device_name(0)}")
|
| log_vram("startup")
|
|
|
| dataset = MeshAIGold5KDataset(token=HF_TOKEN)
|
| train_idx, val_idx = split_train_val_indices(dataset, val_ratio)
|
| train_set = Subset(dataset, train_idx)
|
| val_set = Subset(dataset, val_idx)
|
|
|
| train_loader = DataLoader(
|
| train_set,
|
| batch_size=4,
|
| shuffle=True,
|
| pin_memory=device == "cuda",
|
| collate_fn=collate_batch,
|
| )
|
| val_loader = DataLoader(
|
| val_set,
|
| batch_size=4,
|
| shuffle=False,
|
| pin_memory=device == "cuda",
|
| collate_fn=collate_batch,
|
| )
|
| _log(f"Veri motoru: {len(train_set)} train + {len(val_set)} val nesne.")
|
|
|
| geometry = LoRAAdapter("trellis_geometry").to(device=device, dtype=dtype)
|
| texture = LoRAAdapter("hunyuan_pbr").to(device=device, dtype=dtype)
|
|
|
| unfreeze_all_layers(geometry)
|
| unfreeze_all_layers(texture)
|
|
|
| latest_ckpt = CHECKPOINT_DIR / "latest_model.pt"
|
| global_step = 0
|
| if resume_from and resume_from.exists():
|
| state = torch.load(resume_from, map_location=device, weights_only=True)
|
| if state.get("geometry_adapter") is not None:
|
| geometry.load_state_dict(state["geometry_adapter"])
|
| if state.get("texture_adapter") is not None:
|
| texture.load_state_dict(state["texture_adapter"])
|
| global_step = int(state.get("global_step", 0))
|
| _log(f"Resume modu: {resume_from} (step {global_step})")
|
|
|
| geom_opt = torch.optim.AdamW(geometry.parameters(), lr=5e-5, fused=False)
|
| tex_opt = torch.optim.AdamW(texture.parameters(), lr=5e-5, fused=False)
|
|
|
| def _train_step(adapter, optimizer, batch, loss_name: str) -> float | None:
|
| optimizer.zero_grad(set_to_none=True)
|
| loss = run_batch_step(adapter, batch, device)
|
| if not torch.isfinite(loss):
|
| if monitor.nan_skips < 3:
|
| _log(f"NaN/Inf {loss_name} — batch atlandi.")
|
| monitor.note_nan_skip(loss_name)
|
| optimizer.zero_grad(set_to_none=True)
|
| return None
|
| loss.backward()
|
| torch.nn.utils.clip_grad_norm_(adapter.parameters(), max_norm=1.0)
|
| optimizer.step()
|
| return float(loss.item())
|
|
|
| for epoch in range(1, epochs + 1):
|
| _log(f"--- Epoch {epoch}/{epochs} Baslatildi ---")
|
|
|
| log_vram("before_geometry")
|
| _log(">> Katman 1-3: Microsoft TRELLIS geometrisi fine-tune ediliyor...")
|
| geometry.train()
|
| geom_losses: list[float] = []
|
| for batch in train_loader:
|
| loss_val = _train_step(geometry, geom_opt, batch, "geometry")
|
| if loss_val is None:
|
| continue
|
| global_step += 1
|
| geom_losses.append(loss_val)
|
| monitor.note_step(global_step, "geometry", loss_val)
|
|
|
| if global_step % validation_every == 0:
|
| val_geom = evaluate_adapter(geometry, val_loader, device)
|
| val_tex = evaluate_adapter(texture, val_loader, device)
|
| monitor.note_validation(global_step, val_geom, val_tex)
|
| log_vram(f"validation_step_{global_step}")
|
| step_dir = export_validation_samples(
|
| global_step, val_loader, geometry, texture, device, val_geom, val_tex
|
| )
|
| maybe_upload_validation_to_hf(step_dir, HF_TOKEN or "")
|
|
|
| log_vram("geometry_completed")
|
|
|
| clear_gpu_cache()
|
|
|
| log_vram("before_texture")
|
| _log(">> Katman 4-5: Tencent Hunyuan3D PBR doku motoru fine-tune ediliyor...")
|
| texture.train()
|
| tex_losses: list[float] = []
|
| for batch in train_loader:
|
| loss_val = _train_step(texture, tex_opt, batch, "texture")
|
| if loss_val is None:
|
| continue
|
| global_step += 1
|
| tex_losses.append(loss_val)
|
| monitor.note_step(global_step, "texture", loss_val)
|
|
|
| if global_step % validation_every == 0:
|
| val_geom = evaluate_adapter(geometry, val_loader, device)
|
| val_tex = evaluate_adapter(texture, val_loader, device)
|
| monitor.note_validation(global_step, val_geom, val_tex)
|
| log_vram(f"validation_step_{global_step}")
|
| step_dir = export_validation_samples(
|
| global_step, val_loader, geometry, texture, device, val_geom, val_tex
|
| )
|
| maybe_upload_validation_to_hf(step_dir, HF_TOKEN or "")
|
|
|
| log_vram("texture_completed")
|
|
|
| val_geom = evaluate_adapter(geometry, val_loader, device)
|
| val_tex = evaluate_adapter(texture, val_loader, device)
|
| geom_mean = float(np.mean(geom_losses)) if geom_losses else float("nan")
|
| tex_mean = float(np.mean(tex_losses)) if tex_losses else float("nan")
|
| monitor.note_epoch_end(epoch, epochs, geom_mean, tex_mean, val_geom, val_tex)
|
|
|
| save_checkpoint(latest_ckpt, epoch, global_step, geometry, texture)
|
| epoch_val_dir = OUTPUT_DIR / f"epoch_{epoch:03d}"
|
| export_validation_samples(
|
| global_step, val_loader, geometry, texture, device, val_geom, val_tex
|
| )
|
| latest_step_dir = sorted(OUTPUT_DIR.glob("step_*"))[-1] if list(OUTPUT_DIR.glob("step_*")) else None
|
| if latest_step_dir:
|
| import shutil
|
|
|
| if epoch_val_dir.exists():
|
| shutil.rmtree(epoch_val_dir)
|
| shutil.copytree(latest_step_dir, epoch_val_dir)
|
|
|
| clear_gpu_cache()
|
| _log(f"Epoch {epoch} tamamlandi. Guncel durum bulut hafizasina alindi.")
|
| log_vram(f"epoch_{epoch}_done")
|
|
|
| monitor.finish(ok=monitor.nan_skips < len(train_set))
|
| _log("Egitim tamamlandi.")
|
| _log(f"Ozet log: {PROGRESS_LOG}")
|
| _log(f"Son durum: {STATUS_JSON}")
|
| _log(f"Gorsel ornekler: {OUTPUT_DIR}")
|
|
|
|
|
| def parse_args() -> argparse.Namespace:
|
| parser = argparse.ArgumentParser(description="MeshAI real-data hybrid training pipeline")
|
| parser.add_argument(
|
| "--mode",
|
| choices=("real", "stub", "faz2"),
|
| default="real",
|
| help="real=hibrit kopru; faz2=TRELLIS/Hunyuan LoRA; stub=legacy",
|
| )
|
| parser.add_argument("--epochs", type=int, default=5)
|
| parser.add_argument("--resume_from", type=str, default="")
|
| parser.add_argument("--validation-every", type=int, default=500, help="Kac stepte bir validation")
|
| parser.add_argument("--checkpoint-every", type=int, default=100, help="Kac stepte bir local+HF checkpoint")
|
| parser.add_argument("--val-split", type=float, default=0.1, help="Validation orani (sabit UID yoksa)")
|
| parser.add_argument("--limit", type=int, default=0, help="Smoke: max obje sayisi (0=tum)")
|
| parser.add_argument("--lora-rank", type=int, default=8, help="Faz2 LoRA rank")
|
| parser.add_argument(
|
| "--hf-preprocessed-repo",
|
| default=os.getenv("HF_PREPROCESSED_REPO", "HayrettinIscan/MeshAI-Preprocessed-4K"),
|
| )
|
| parser.add_argument("--data-root", default="", help="Yerel preprocessed kok (opsiyonel)")
|
| return parser.parse_args()
|
|
|
|
|
| if __name__ == "__main__":
|
| if not HF_TOKEN:
|
| _log("HATA: HF_TOKEN eksik! Bulut dogrulamasi olmadan egitim baslatilamaz.")
|
| sys.exit(1)
|
|
|
| args = parse_args()
|
| resume = Path(args.resume_from) if args.resume_from else None
|
| limit = args.limit if args.limit > 0 else None
|
| data_root = Path(args.data_root) if args.data_root else None
|
|
|
| if args.mode == "faz2":
|
| from meshai_train.faz2_engine import start_faz2_training
|
|
|
| monitor = TrainMonitor()
|
| monitor.last["version"] = "v5.0-faz2"
|
| monitor._save_status()
|
| start_faz2_training(
|
| monitor=monitor,
|
| checkpoint_dir=CHECKPOINT_DIR,
|
| output_dir=OUTPUT_DIR,
|
| token=HF_TOKEN,
|
| epochs=args.epochs,
|
| resume_from=resume,
|
| validation_every=args.validation_every,
|
| val_ratio=args.val_split,
|
| limit=limit,
|
| hf_repo=args.hf_preprocessed_repo,
|
| data_root=data_root,
|
| log_fn=_log,
|
| log_vram_fn=log_vram,
|
| clear_gpu_fn=clear_gpu_cache,
|
| load_val_uids_fn=load_validation_uids,
|
| checkpoint_every=args.checkpoint_every,
|
| lora_rank=args.lora_rank,
|
| )
|
| elif args.mode == "real":
|
| from meshai_train.engine import start_real_training
|
|
|
| monitor = TrainMonitor()
|
| monitor.last["version"] = "v4.0-real"
|
| monitor._save_status()
|
| start_real_training(
|
| monitor=monitor,
|
| checkpoint_dir=CHECKPOINT_DIR,
|
| output_dir=OUTPUT_DIR,
|
| token=HF_TOKEN,
|
| epochs=args.epochs,
|
| resume_from=resume,
|
| validation_every=args.validation_every,
|
| val_ratio=args.val_split,
|
| limit=limit,
|
| hf_repo=args.hf_preprocessed_repo,
|
| data_root=data_root,
|
| log_fn=_log,
|
| log_vram_fn=log_vram,
|
| clear_gpu_fn=clear_gpu_cache,
|
| load_val_uids_fn=load_validation_uids,
|
| checkpoint_every=args.checkpoint_every,
|
| )
|
| else:
|
| start_training(
|
| epochs=args.epochs,
|
| resume_from=resume,
|
| validation_every=args.validation_every,
|
| val_ratio=args.val_split,
|
| )
|
|
|
