PROCITY/pipeline/gen_props.py

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

Two local stages, both on the M3 Ultra GPU, zero cloud cost:
  1. concept image  flux_local (FLUX.2-klein-4B)  text->object photo   ~6s
  2. mesh           trellis_mac (TRELLIS.2-4B)     image->GLB+PBR       ~3-5 min (1024_cascade)
then pipeline/normalize.py (house GLB law) + thumbnail + manifest.

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   # stage 2 for one (concept must exist)
    python3 pipeline/gen_props.py --mesh                # stage 2 for every prop with a concept, no glb
    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")
TRELLIS_PY = os.path.join(MB, "vendor/trellis-mac/.venv/bin/python")
TRELLIS_RUN = os.path.join(MB, "server/operators/trellis_mac/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 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):
    """concept image -> GLB. TRELLIS 1024_cascade by default; sf3d for a fast draft."""
    name = a["name"]
    cpath = concept_path(name)
    if not os.path.exists(cpath):
        raise RuntimeError(f"no concept image for {name} — run --concepts first")
    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
    else:
        params = json.dumps({"pipeline_type": "1024_cascade", "texture_size": 1024, "seed": 0})
        cmd = [TRELLIS_PY, TRELLIS_RUN, "--input", cpath, "--outdir", outdir, "--params", params]
        tmo = 1200
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=tmo)
    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:
        raise RuntimeError((r.stderr or r.stdout or "no glb")[-240:])
    dst = glb_path(name)
    shutil.move(sorted(glbs)[0], 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
    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 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)
                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.")