File size: 2,646 Bytes
29665a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #!/usr/bin/env python3
"""Smoke Faz3 TRELLIS forward LoRA then optional Paint LoRA. Uses HF Preprocessed-4K."""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--paint", action="store_true", help="Ayrica faz3_paint smoke")
parser.add_argument("--limit", type=int, default=4)
parser.add_argument("--epochs", type=int, default=1)
parser.add_argument("--lora-rank", type=int, default=8)
args = parser.parse_args()
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
if not token:
print("HF_TOKEN gerekli")
return 1
os.environ.setdefault("ATTN_BACKEND", "xformers")
os.environ.setdefault("SPARSE_ATTN_BACKEND", "xformers")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# Re-enter train_pipeline CLI
cmd_geo = [
sys.executable,
str(ROOT / "train_pipeline.py"),
"--mode",
"faz3",
"--epochs",
str(args.epochs),
"--limit",
str(args.limit),
"--lora-rank",
str(args.lora_rank),
"--checkpoint-every",
"20",
"--validation-every",
"50",
]
print("RUN", " ".join(cmd_geo))
import subprocess
env = os.environ.copy()
env["HF_TOKEN"] = token
r = subprocess.call(cmd_geo, cwd=str(ROOT), env=env)
if r != 0:
return r
ckpt = ROOT / "checkpoints" / "latest_faz3_lora.pt"
if not ckpt.is_file():
print("FAIL: latest_faz3_lora.pt yok")
return 1
print("OK geo ckpt", ckpt, ckpt.stat().st_size)
if args.paint:
cmd_p = [
sys.executable,
str(ROOT / "train_pipeline.py"),
"--mode",
"faz3_paint",
"--epochs",
str(args.epochs),
"--limit",
str(min(2, args.limit)),
"--lora-rank",
str(args.lora_rank),
"--checkpoint-every",
"5",
]
print("RUN", " ".join(cmd_p))
r2 = subprocess.call(cmd_p, cwd=str(ROOT), env=env)
if r2 != 0:
return r2
ckpt2 = ROOT / "checkpoints" / "latest_faz3_paint_lora.pt"
print("OK paint ckpt", ckpt2, ckpt2.stat().st_size if ckpt2.is_file() else None)
print("SMOKE OK")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|