PROCITY/pipeline/gen_props.py

#!/usr/bin/env python3
"""PROCITY hero props — free, on-device (round-3 Fable path: MODELBEAST, not fal.ai).

Three local stages, all on the M3 Ultra GPU, zero cloud cost:
  1. concept image  flux_local (FLUX.2-klein-4B)     text->object photo    ~6-9s
  2. cutout         bg_remove_local (RMBG-2.0)       photo->RGBA subject   ~5-10s
  3. mesh           trellis2_mlx (TRELLIS.2-4B MLX)  image->GLB+PBR        ~2-3 min (1024_cascade)
then pipeline/normalize.py (house GLB law) + thumbnail + manifest.

R38 PIPELINE FIX — the operator swap (docs/REVIEW_CITY_PATTERNS.md §2, measured on the farm):
this script used to hardcode `trellis_mac` (p50 245.4s, 33 jobs/7d). The farm's workhorse is
**`trellis2_mlx`** (384 jobs/7d, p50 134.6s, better PBR bake) — a ~2x throughput gain. The two
operators are NOT drop-in on output: trellis_mac writes one `*.glb`, trellis2_mlx writes BOTH
`candidate.glb` (raw, untextured) and `candidate_pbr.glb` (the baked one) into the same tree,
so a naive `sorted(rglob('*.glb'))[0]` silently picks the untextured mesh. We read `result.json`
(the operator's own declared output) and fall back to `candidate_pbr.glb` by name. `--legacy-trellis`
keeps the old operator available for an A/B.

Prop list + prompts come from pipeline/meshgod_batch.json (the same gap list that was going to
fal.ai). cash-register is skipped — vintage-cash-register.glb already lives on the 3GOD depot.

    python3 pipeline/gen_props.py --dry-run             # list props + what's done, no GPU
    python3 pipeline/gen_props.py --concepts            # stage 1 for all -> .genprops/<name>.png
    python3 pipeline/gen_props.py --concepts --only glass-case
    python3 pipeline/gen_props.py --mesh --only glass-case   # stages 2+3 for one (concept must exist)
    python3 pipeline/gen_props.py --mesh                # stages 2+3 for every prop with a concept, no glb
    python3 pipeline/gen_props.py --mesh --no-cutout    # skip RMBG (concept already has alpha)
    python3 pipeline/gen_props.py --mesh --legacy-trellis    # the old trellis_mac path (A/B only)
    python3 pipeline/gen_props.py --draft --only glass-case # sf3d seconds-fast draft (iterate concept)

Resumable: skips a stage whose output already exists in .genprops/. Eyeball .genprops/*.png (concepts)
and the rendered GLBs before normalizing. Everything here is OPTIONAL at runtime — every prop has an
in-engine primitive fallback (see MESHGOD_BATCH.md _primitives_instead), so nothing blocks on it.
"""
import json, os, sys, subprocess, shutil, time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.path.join(ROOT, "pipeline", ".genprops")
BATCH = os.path.join(ROOT, "pipeline", "meshgod_batch.json")
os.makedirs(RAW, exist_ok=True)

MB = os.environ.get("MB_HOME", os.path.expanduser("~/Documents/MODELBEAST"))
FLUX_PY = os.path.join(MB, "venvs/mflux/bin/python")
FLUX_RUN = os.path.join(MB, "server/operators/flux_local/run.py")
# R38: trellis2_mlx is the farm's workhorse (p50 134.6s vs trellis_mac's 245.4s). Its operator
# run.py is stdlib-only and resolves the fork's own .venv internally, so the outer interpreter
# just has to exist; we use the fork venv for symmetry with the legacy line below.
TRELLIS2_PY = os.path.join(MB, "vendor/trellis2-mlx/.venv/bin/python")
TRELLIS2_RUN = os.path.join(MB, "server/operators/trellis2_mlx/run.py")
TRELLIS_PY = os.path.join(MB, "vendor/trellis-mac/.venv/bin/python")     # legacy (--legacy-trellis)
TRELLIS_RUN = os.path.join(MB, "server/operators/trellis_mac/run.py")
RMBG_PY = os.path.join(MB, "venvs/rmbg/bin/python")
RMBG_RUN = os.path.join(MB, "server/operators/bg_remove_local/run.py")
SF3D_PY = os.path.join(MB, "venvs/sf3d/bin/python")
SF3D_RUN = os.path.join(MB, "server/operators/sf3d/run.py")

# Clean product-shot framing so TRELLIS/sf3d (which segment the subject) get a whole, isolated
# object on a plain backdrop — NOT the faded op-shop facade look. Period-appropriate but neutral.
CONCEPT = ("Clean studio product photograph of {v}. Single object, centred, entire object in "
           "frame, three-quarter view, plain seamless pale-grey backdrop, soft even lighting, "
           "sharp focus, no text, no people, no other objects, 1990s Australian period-correct.")


def props():
    data = json.load(open(BATCH))
    out = []
    for a in data["assets"]:
        # cash-register: reuse the depot's vintage-cash-register.glb (round-3 reuse check)
        if a["name"] == "cash-register":
            continue
        out.append(a)
    return out


def concept_path(name):
    return os.path.join(RAW, f"{name}.png")


def glb_path(name):
    return os.path.join(RAW, f"{name}.glb")


def cutout_path(name):
    return os.path.join(RAW, f"{name}.cut.png")


def gen_cutout(name):
    """concept photo -> RGBA subject cutout (bg_remove_local / RMBG-2.0). The mesh operators
    segment internally, but feeding a clean alpha measurably improves thin-feature retention and
    it is the path the R38 review measured. Fail-soft: caller falls back to the raw concept."""
    src = concept_path(name)
    outdir = os.path.join(RAW, "_x_" + name)
    os.makedirs(outdir, exist_ok=True)
    params = json.dumps({"resolution": 1024, "background": "transparent"})
    r = subprocess.run([RMBG_PY, RMBG_RUN, "--input", src, "--outdir", outdir, "--params", params],
                       capture_output=True, text=True, timeout=600)
    cut = os.path.join(outdir, os.path.splitext(os.path.basename(src))[0] + "_cutout.png")
    if not os.path.exists(cut):
        shutil.rmtree(outdir, ignore_errors=True)
        raise RuntimeError((r.stderr or r.stdout or "no cutout")[-200:])
    dst = cutout_path(name)
    shutil.move(cut, dst)
    shutil.rmtree(outdir, ignore_errors=True)
    return dst


def _pick_glb(outdir):
    """The operator-output contract, made explicit. trellis2_mlx emits BOTH candidate.glb (raw)
    and candidate_pbr.glb (baked) — alphabetical order picks the WRONG one. Trust result.json
    first (each operator declares its own output there), then the _pbr name, then anything."""
    rj = os.path.join(outdir, "result.json")
    if os.path.exists(rj):
        try:
            outs = json.load(open(rj)).get("outputs") or []
            for o in outs:
                p = o.get("path", "")
                p = p if os.path.isabs(p) else os.path.join(outdir, p)
                if p.endswith(".glb") and os.path.exists(p):
                    return p
        except Exception:
            pass
    glbs = []
    for dp, _, fs in os.walk(outdir):
        glbs += [os.path.join(dp, f) for f in fs if f.endswith(".glb")]
    if not glbs:
        return None
    pbr = [g for g in glbs if os.path.basename(g) == "candidate_pbr.glb"]
    return (pbr or sorted(glbs))[0]


def gen_concept(a, seed=3):
    name, prompt = a["name"], a["prompt"]
    outdir = os.path.join(RAW, "_c_" + name)
    os.makedirs(outdir, exist_ok=True)
    params = json.dumps({"prompt": CONCEPT.format(v=prompt), "model": "flux2-klein-4b",
                         "steps": 4, "width": 1024, "height": 1024, "seed": seed})
    r = subprocess.run([FLUX_PY, FLUX_RUN, "--outdir", outdir, "--params", params],
                       capture_output=True, text=True, timeout=300)
    pngs = [p for p in os.listdir(outdir) if p.endswith(".png")]
    if not pngs:
        raise RuntimeError((r.stderr or r.stdout or "no png")[-200:])
    shutil.move(os.path.join(outdir, pngs[0]), concept_path(name))
    shutil.rmtree(outdir, ignore_errors=True)
    return concept_path(name)


def gen_mesh(a, draft=False, legacy=False, cutout=True):
    """concept image -> GLB. trellis2_mlx 1024_cascade by default (R38 swap); sf3d for a fast
    draft; --legacy-trellis for the old trellis_mac operator."""
    name = a["name"]
    cpath = concept_path(name)
    if not os.path.exists(cpath):
        raise RuntimeError(f"no concept image for {name} — run --concepts first")
    if cutout and not draft:
        if os.path.exists(cutout_path(name)):
            cpath = cutout_path(name)
        else:
            try:
                cpath = gen_cutout(name)
            except Exception as e:                       # fail-soft: the mesh op segments anyway
                print(f"  [cutout] {name} skipped: {str(e)[:90]}")
    outdir = os.path.join(RAW, "_m_" + name)
    os.makedirs(outdir, exist_ok=True)
    if draft:
        params = json.dumps({"device": "mps", "texture_resolution": 1024, "foreground_ratio": 0.85})
        cmd = [SF3D_PY, SF3D_RUN, "--input", cpath, "--outdir", outdir, "--params", params]
        tmo = 600
    elif legacy:
        params = json.dumps({"pipeline_type": "1024_cascade", "texture_size": 1024, "seed": 0})
        cmd = [TRELLIS_PY, TRELLIS_RUN, "--input", cpath, "--outdir", outdir, "--params", params]
        tmo = 1800
    else:
        # trellis2_mlx params_schema (manifest.json, verified R38): pipeline_type / seed / steps /
        # baker / texture_size / max_bake_faces / alpha_mode. metal baker = UV PBR (meshgod-grade);
        # 500k bake cap = fal parity and M3-Ultra-verified; opaque alpha kills the speckle veil.
        params = json.dumps({"pipeline_type": "1024_cascade", "texture_size": 1024, "seed": 0,
                             "baker": "metal", "max_bake_faces": 500000, "alpha_mode": "opaque"})
        cmd = [TRELLIS2_PY, TRELLIS2_RUN, "--input", cpath, "--outdir", outdir, "--params", params]
        tmo = 1800
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=tmo)
    src = _pick_glb(outdir)
    if not src:
        raise RuntimeError((r.stderr or r.stdout or "no glb")[-240:])
    dst = glb_path(name)
    shutil.move(src, dst)
    shutil.rmtree(outdir, ignore_errors=True)
    return dst


if __name__ == "__main__":
    if "--batch" in sys.argv:                       # override the asset list (e.g. gig instruments)
        BATCH = os.path.abspath(sys.argv[sys.argv.index("--batch") + 1])
    only = sys.argv[sys.argv.index("--only") + 1] if "--only" in sys.argv else ""
    draft = "--draft" in sys.argv
    legacy = "--legacy-trellis" in sys.argv
    cutout = "--no-cutout" not in sys.argv
    todo = [a for a in props() if not only or only in a["name"]]

    if "--dry-run" in sys.argv:
        print(f"{len(todo)} props (cash-register reused from depot):")
        for a in todo:
            c = "concept✓" if os.path.exists(concept_path(a["name"])) else "concept·"
            g = "glb✓" if os.path.exists(glb_path(a["name"])) else "glb·"
            print(f"  {a['name']:16s} {c} {g}  h={a['height_m']}m  for={a['for']}")
        sys.exit(0)

    do_concepts = "--concepts" in sys.argv or (not "--mesh" in sys.argv and not draft)
    do_mesh = "--mesh" in sys.argv or draft

    if do_concepts:
        for a in todo:
            if os.path.exists(concept_path(a["name"])):
                print(f"[concept] {a['name']} — exists, skip"); continue
            t0 = time.time()
            try:
                gen_concept(a)
                print(f"[concept] {a['name']} — {round(time.time()-t0,1)}s")
            except Exception as e:
                print(f"[concept] {a['name']} FAILED: {e}")

    if do_mesh:
        tool = ("sf3d-draft" if draft else
                ("trellis_mac 1024_cascade" if legacy else "trellis2_mlx 1024_cascade"))
        for a in todo:
            if not draft and os.path.exists(glb_path(a["name"])):
                print(f"[mesh] {a['name']} — glb exists, skip"); continue
            t0 = time.time()
            try:
                dst = gen_mesh(a, draft=draft, legacy=legacy, cutout=cutout)
                kb = os.path.getsize(dst) // 1024
                print(f"[mesh:{tool}] {a['name']} — {round(time.time()-t0,1)}s ({kb}KB)")
            except Exception as e:
                print(f"[mesh:{tool}] {a['name']} FAILED: {e}")
    print("done. Eyeball .genprops/*.png (concepts) + rendered GLBs before normalize.")