AIGame/Tools/output/run_townsquare_pipeline.py

"""TownSquare 分件：图生3D → 减面烘焙 → 安装到 Bundles/Scenes/Town/TownSquare/{name}/"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT / "Tools") not in sys.path:
    sys.path.insert(0, str(ROOT / "Tools"))
PARTS = [
    ("ground", 800),
    ("shop", 1800),
    ("tree", 1500),
    ("lamp", 500),
    ("board", 800),
    ("fountain", 1800),
    ("bench", 700),
    ("bush", 500),
]
SCENE_TYPE = "Town"
SCENE_ID = "TownSquare"
TEX_SIZE = 1024


def _decode(raw: bytes | None) -> str:
    if not raw:
        return ""
    return raw.decode("utf-8", errors="replace")


def run(cmd: list[str]) -> dict:
    print(f"RUN {' '.join(cmd)}", flush=True)
    env = dict(os.environ)
    env["PYTHONIOENCODING"] = "utf-8"
    proc = subprocess.run(cmd, cwd=ROOT, capture_output=True, env=env)
    out = _decode(proc.stdout).strip()
    err = _decode(proc.stderr).strip()
    if proc.returncode != 0:
        raise RuntimeError(
            f"exit {proc.returncode}\nstdout:\n{out[-4000:]}\nstderr:\n{err[-4000:]}"
        )
    last = None
    for line in out.splitlines():
        line = line.strip()
        if line.startswith("{") and line.endswith("}"):
            try:
                last = json.loads(line)
            except json.JSONDecodeError:
                continue
    if last is None and out.startswith("{"):
        try:
            last = json.loads(out)
        except json.JSONDecodeError:
            last = None
    return last or {"ok": True, "raw": out[-2000:]}


def is_valid_glb(path: Path) -> bool:
    from hunyuan_3d.glb_check import glb_textures_ok

    ok, _ = glb_textures_ok(path)
    return ok


def pick(files: list[Path], pred) -> Path | None:
    for path in files:
        if path.is_file() and pred(path.name.lower()):
            return path
    return None


def install_part(name: str, xpbr_dir: Path, concept: Path) -> Path:
    dest = ROOT / "AIClient/Assets/Bundles/Scenes" / SCENE_TYPE / SCENE_ID / name
    dest.mkdir(parents=True, exist_ok=True)
    files = list(xpbr_dir.iterdir())
    fbx = pick(files, lambda n: n.endswith(".fbx"))
    base = pick(files, lambda n: n.endswith("_base.png") or n.endswith("_basemap.png"))
    mix = pick(files, lambda n: n.endswith("_mix.png") or n.endswith("_mixmap.png"))
    mat = pick(files, lambda n: n.endswith(".mat"))
    if not fbx or not base:
        raise FileNotFoundError(f"XPbr 不完整: {xpbr_dir} files={[p.name for p in files]}")
    mapping = {
        fbx: dest / f"{name}.fbx",
        base: dest / f"{name}_base.png",
        mix: dest / f"{name}_mix.png" if mix else None,
        mat: dest / f"{name}.mat" if mat else None,
    }
    copied = []
    for src, target in mapping.items():
        if src is None or target is None:
            continue
        shutil.copy2(src, target)
        meta = Path(str(src) + ".meta")
        if meta.is_file():
            shutil.copy2(meta, Path(str(target) + ".meta"))
        copied.append(target.name)
    print(json.dumps({"installed": str(dest), "files": copied, "design": str(concept)}, ensure_ascii=False), flush=True)
    return dest


def main() -> int:
    summary = []
    for name, tris in PARTS:
        img = ROOT / "Doc/Design/Scene/TownSquare" / f"{name}.png"
        glb = ROOT / "Tools/output/townsquare_3d" / name / f"{name}.glb"
        bake_out = ROOT / "Tools/output/townsquare_game" / name
        glb.parent.mkdir(parents=True, exist_ok=True)
        print(json.dumps({"event": "part_start", "name": name, "tris": tris}, ensure_ascii=False), flush=True)

        if not is_valid_glb(glb):
            if glb.is_file():
                print(
                    json.dumps({"event": "hunyuan_regen", "name": name, "reason": "invalid_glb"}, ensure_ascii=False),
                    flush=True,
                )
                glb.unlink()
            hunyuan = run(
                [
                    sys.executable,
                    str(ROOT / "Tools/hunyuan_3d_generate.py"),
                    "-i",
                    str(img),
                    "-o",
                    str(glb),
                    "--faces",
                    "50k",
                    "--result-only",
                    "--quiet",
                ]
            )
            if hunyuan.get("ok") is False or hunyuan.get("status") not in (None, "succeeded"):
                raise RuntimeError(f"hunyuan failed {name}: {hunyuan}")
            if not is_valid_glb(glb):
                raise RuntimeError(f"hunyuan produced invalid glb: {glb}")
        else:
            print(
                json.dumps({"event": "hunyuan_skip", "name": name}, ensure_ascii=False),
                flush=True,
            )

        xpbr_dir = bake_out / f"{name}_xpbr"
        bake: dict = {}
        has_xpbr = xpbr_dir.is_dir() and any(xpbr_dir.glob("*.fbx"))
        if has_xpbr:
            print(
                json.dumps({"event": "bake_skip", "name": name}, ensure_ascii=False),
                flush=True,
            )
        else:
            bake = run(
                [
                    sys.executable,
                    str(ROOT / "Tools/model_tools/model_lowpoly_bake.py"),
                    str(glb),
                    "-o",
                    str(bake_out),
                    "--tris",
                    str(tris),
                    "--size",
                    str(TEX_SIZE),
                    "--result-only",
                    "--quiet",
                ]
            )
            if bake.get("xpbr_dir"):
                xpbr_dir = Path(bake["xpbr_dir"])
        dest = install_part(name, xpbr_dir, img)
        rec = {
            "name": name,
            "tris_target": tris,
            "tris_after": bake.get("tris_after"),
            "glb": str(glb.relative_to(ROOT)).replace("\\", "/"),
            "dest": str(dest.relative_to(ROOT)).replace("\\", "/"),
        }
        summary.append(rec)
        print(json.dumps({"event": "part_done", **rec}, ensure_ascii=False), flush=True)

    out = ROOT / "Tools/output/townsquare_pipeline_summary.json"
    out.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"ok": True, "summary": summary}, ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False), flush=True)
        raise