Blame for projtex/cli.py

"""Command line interface. Every command prints one JSON document on stdout (logs go to stderr).

    projtex init PROJ mesh.glb --mode character
    projtex config PROJ --prompt "..." --backend qwen-local
    projtex plan PROJ
    projtex generate PROJ
    projtex bake PROJ --resolution 4096
    projtex inspect PROJ
    projtex retouch PROJ --focus 0.1,1.6,0.05,0.12 --mask low_quality --mode repair
    projtex render PROJ --compare-source
    projtex run mesh.glb OUT --prompt "..."            (all of the above in one go)
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import traceback
from typing import Any


def _parse_value(s: str) -> Any:
    try:
        return json.loads(s)
    except (json.JSONDecodeError, TypeError):
        return s


def _kv(items: list[str] | None) -> dict:
    out = {}
    for it in items or []:
        if "=" not in it:
            raise SystemExit(f"expected key=value, got {it!r}")
        k, v = it.split("=", 1)
        out[k.strip()] = _parse_value(v)
    return out


def _floats(s: str | None) -> list[float] | None:
    if s is None:
        return None
    return [float(x) for x in s.split(",")]


def _focus(s: str | None):
    if s is None:
        return None
    if s.strip().startswith("{"):
        return json.loads(s)
    if s in ("auto", "head", "face", "none"):
        return None if s == "none" else s
    v = _floats(s)
    if len(v) == 4:
        return {"point": v[:3], "radius": v[3]}
    if len(v) == 7:
        return {"point": v[:3], "radius": v[3], "direction": v[4:]}
    raise SystemExit("focus must be x,y,z,r[,dx,dy,dz] or JSON")


def _mask(s: str | None):
    if s is None:
        return None
    s = s.strip()
    if s.startswith("{") or s.startswith("["):
        return json.loads(s)
    return s


def _backend(args) -> Any:
    if not getattr(args, "backend", None) and not getattr(args, "backend_opt", None):
        return None
    d = _kv(getattr(args, "backend_opt", None))
    if args.backend:
        d["name"] = args.backend
    return d


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(prog="projtex", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("-v", "--verbose", action="store_true", help="INFO logs on stderr")
    sp = ap.add_subparsers(dest="cmd", required=True)

    def proj(p):
        p.add_argument("project", help="project directory")
        return p

    def backend_args(p):
        p.add_argument("--backend", help="qwen-turbo (default) | qwen-space | fibo-space | qwen-local | hf | files | oracle | mock")
        p.add_argument("--backend-opt", action="append", metavar="K=V",
                       help="backend option, e.g. prompt_extend=true or model=Qwen/Qwen-Image-2.1 (repeatable)")

    p = sp.add_parser("init", help="import a mesh and create a project")
    p.add_argument("project")
    p.add_argument("mesh", help=".glb/.gltf/.obj/.ply")
    p.add_argument("--mode", default="object", choices=["object", "character", "environment"])
    p.add_argument("--uv", default="auto", choices=["auto", "keep", "pack", "regen"])
    p.add_argument("--up", choices=["auto", "y", "z", "-z", "x"],
                   help="up axis of the asset (default y; auto in character mode)")
    p.add_argument("--texture-size", type=int)
    p.add_argument("--view-size", type=int)
    p.add_argument("--overwrite", action="store_true")

    p = proj(sp.add_parser("config", help="update settings"))
    p.add_argument("--prompt")
    p.add_argument("--subject", help="what the mesh is, e.g. 'a woman', 'a kitchen'")
    p.add_argument("--negative-prompt")
    p.add_argument("--lighting", choices=["flat", "soft", "natural"])
    p.add_argument("--raw-prompt", type=_parse_value, help="true: send the prompt verbatim")
    p.add_argument("--style-reference", help="global style reference image")
    p.add_argument("--preset", help="append a material preset: ceramic bronze scifi stone abandoned workshop wood concrete")
    p.add_argument("--texture-size", type=int)
    p.add_argument("--view-size", type=int)
    p.add_argument("--preview-texture-size", type=int)
    p.add_argument("--seed", type=int)
    p.add_argument("--skymap", help="equirectangular background image for environments")
    p.add_argument("--front-azimuth", type=_parse_value,
                   help="degrees to rotate the 'front' camera if the asset does not face +Z, or \"auto\"")
    p.add_argument("--fusion", action="append", metavar="K=V", help="fusion parameter, e.g. sharpness=6 (repeatable)")
    p.add_argument("--set", action="append", metavar="K=V", help="any other setting as JSON value (repeatable)")
    backend_args(p)

    p = proj(sp.add_parser("plan", help="place cameras"))
    p.add_argument("--preset", help="object: minimal | balanced | dense | greedy; environment: probes | probes-greedy")
    p.add_argument("--views", type=int, help="number of views for --preset greedy")
    p.add_argument("--focus", help="auto | none | x,y,z,r : close-up views (default auto in character mode)")
    p.add_argument("--no-gaps", action="store_true", help="do not add gap-filling views")
    p.add_argument("--max-gap-views", type=int)
    p.add_argument("--probes", type=int, help="environment: number of 360 probes (default 6)")
    p.add_argument("--eye-height", type=float)
    p.add_argument("--levels", type=_parse_value, default="lowest",
                   help="environment floor levels for probes: lowest (default) | all | a height")
    p.add_argument("--region", default="auto", choices=["auto", "enclosed", "all"],
                   help="environment: keep only enclosed walkable space (auto drops ground outside the building)")
    p.add_argument("--fov", type=float)
    p.add_argument("--append", action="store_true", help="keep existing views")

    proj(sp.add_parser("status", help="project summary"))

    p = proj(sp.add_parser("view", help="change one view"))
    p.add_argument("name")
    p.add_argument("--enable", action="store_true")
    p.add_argument("--disable", action="store_true")
    p.add_argument("--priority", type=float)
    p.add_argument("--reference", help="per-view reference image (same angle)")
    p.add_argument("--prompt-extra")
    p.add_argument("--seed", type=int)
    p.add_argument("--remove", action="store_true")

    p = proj(sp.add_parser("add-view", help="add a custom camera"))
    p.add_argument("name")
    p.add_argument("--eye")
    p.add_argument("--target")
    p.add_argument("--fov", type=float, default=40.0)
    p.add_argument("--focus", help="x,y,z,r[,dx,dy,dz]: frame a sphere")
    p.add_argument("--priority", type=float, default=1.0)
    p.add_argument("--index", type=int, help="insert position in generation order")

    p = proj(sp.add_parser("generate", help="generate & project views (sequential, conditioned on previous ones)"))
    p.add_argument("--views", help="comma-separated subset")
    p.add_argument("--force", action="store_true", help="regenerate views that already have images")
    p.add_argument("--no-bake", action="store_true")
    backend_args(p)

    p = proj(sp.add_parser("fill-gaps", help="add + generate cameras for surface that is still uncovered"))
    p.add_argument("--max-views", type=int, default=3)
    p.add_argument("--rounds", type=int, default=1)
    p.add_argument("--plan-only", action="store_true")
    backend_args(p)

    p = proj(sp.add_parser("condition", help="write the conditioning render + prompt of a view (external generation)"))
    p.add_argument("name")

    p = proj(sp.add_parser("import-view", help="use an external image for a view"))
    p.add_argument("name")
    p.add_argument("image")

    p = proj(sp.add_parser("bake", help="fuse views into the UV texture and export textured.glb"))
    p.add_argument("--resolution", type=int, help="2048 / 4096 / 8192")

    p = proj(sp.add_parser("export", help="export glb/obj"))
    p.add_argument("--format", default="glb", choices=["glb", "obj"])
    p.add_argument("--out")
    p.add_argument("--resolution", type=int)
    p.add_argument("--unlit", action="store_true", help="add KHR_materials_unlit")

    p = proj(sp.add_parser("inspect", help="coverage/quality stats and retouch suggestions"))
    p.add_argument("--max-suggestions", type=int, default=6)

    p = proj(sp.add_parser("retouch", help="local repair / replace / upscale inside a mask"))
    p.add_argument("--prompt", default="")
    p.add_argument("--mode", default="repair", choices=["repair", "replace", "upscale"])
    p.add_argument("--view", help="reuse a view's camera")
    p.add_argument("--focus", help="x,y,z,r[,dx,dy,dz] or JSON {point, radius, direction, fov}")
    p.add_argument("--camera", help='JSON {"eye":[..],"target":[..],"fov":40}')
    p.add_argument("--render-camera", help="reuse a preview camera, e.g. walk01 or orbit03 (from renders/cameras.json)")
    p.add_argument("--suggestion", type=int, help="use suggestion #N from `inspect` (focus + mask + mode)")
    p.add_argument("--mask", help="all | uncovered | low_quality | mask.png | JSON {sphere|box|rect|circle|texels}")
    p.add_argument("--capture-size", type=int, help="render size of the capture (e.g. 256 for the upscale trick)")
    p.add_argument("--output-size", type=int, help="generated size (default view size)")
    p.add_argument("--softness", type=float, default=0.75, help="feather width as a fraction of the mask radius")
    p.add_argument("--grow", type=int, default=0, help="dilate the mask by N pixels")
    p.add_argument("--mark", choices=["prefill", "fill", "outline", "mask", "none"],
                   help="how the region is shown to the model (default: replace=prefill, repair=outline, upscale=none)")
    p.add_argument("--strength", type=float, default=1.0)
    p.add_argument("--seed", type=int)
    p.add_argument("--name")
    p.add_argument("--no-color-match", action="store_true")
    backend_args(p)

    p = proj(sp.add_parser("pick", help="3D point / normal / UV under image pixels of a view or camera"))
    p.add_argument("pixels", help="x,y[;x,y...] in pixels or 0-1 fractions")
    p.add_argument("--view")
    p.add_argument("--camera", help="JSON camera")
    p.add_argument("--render-camera", help="a preview camera from renders/cameras.json (e.g. walk01)")
    p.add_argument("--size", type=int, help="image size the pixel coordinates refer to")

    proj(sp.add_parser("undo", help="disable the last retouch layer"))
    p = proj(sp.add_parser("layer", help="enable/disable a retouch layer"))
    p.add_argument("name")
    p.add_argument("--enable", action="store_true")
    p.add_argument("--disable", action="store_true")
    p.add_argument("--strength", type=float)

    p = proj(sp.add_parser("render", help="preview renders, comparisons, turntable video"))
    p.add_argument("--turntable", type=int, default=8)
    p.add_argument("--size", type=int, default=768)
    p.add_argument("--mode", default="albedo", choices=["albedo", "lit", "clay", "normals"])
    p.add_argument("--elevation", type=float, default=15.0)
    p.add_argument("--views", help="render these view cameras instead of an orbit")
    p.add_argument("--compare-source", action="store_true", help="side by side with the input / original texture")
    p.add_argument("--video", action="store_true")
    p.add_argument("--out-dir", default="renders")
    p.add_argument("--sheet", action="store_true", help="also write a (condition | generated) sheet of all views")

    p = sp.add_parser("run", help="one-shot: init + plan + generate + bake + render")
    p.add_argument("mesh")
    p.add_argument("project")
    p.add_argument("--mode", default="object", choices=["object", "character", "environment"])
    p.add_argument("--prompt", required=True)
    p.add_argument("--subject")
    p.add_argument("--preset")
    p.add_argument("--probes", type=int)
    p.add_argument("--focus")
    p.add_argument("--texture-size", type=int)
    p.add_argument("--view-size", type=int)
    p.add_argument("--uv", default="auto")
    p.add_argument("--up", choices=["auto", "y", "z", "-z", "x"])
    p.add_argument("--style-reference")
    p.add_argument("--lighting", choices=["flat", "soft", "natural"])
    p.add_argument("--overwrite", action="store_true")
    p.add_argument("--auto-retouch", type=int, default=0, help="apply up to N inspect suggestions afterwards")
    backend_args(p)

    sp.add_parser("backends", help="list generation backends")
    sp.add_parser("schema", help="machine-readable description of all commands")
    p = sp.add_parser("mcp", help="run a Model Context Protocol server on stdio")
    return ap


def _schema(ap: argparse.ArgumentParser) -> dict:
    out = {}
    for act in ap._subparsers._group_actions:  # noqa: SLF001
        for name, sub in act.choices.items():
            out[name] = {
                "help": sub.description or next((c.help for c in act._choices_actions if c.dest == name), ""),  # noqa: SLF001
                "args": [{"name": a.dest, "flags": a.option_strings, "help": a.help, "default": a.default,
                          "required": a.required or not a.option_strings,
                          "choices": list(a.choices) if a.choices else None}
                         for a in sub._actions if a.dest != "help"],  # noqa: SLF001
            }
    return out


def run(args, ap) -> Any:
    from .project import Project

    c = args.cmd
    if c == "backends":
        from .backends import BACKENDS

        return {"backends": sorted(BACKENDS), "notes": {
            "qwen-turbo": "default. Qwen-Image-2.1-Turbo (8 steps, ~7 s of GPU per view) through the Space hugging-apps/qwen-qwen-image-2-1-turbo (ZeroGPU quota of your HF login; needs `pip install projtex[spaces]`)",
            "qwen-space": "full Qwen-Image-2.1 (40 steps) through the official Space Qwen/Qwen-Image-2.1 (ZeroGPU quota of your HF login; needs `pip install projtex[spaces]`)",
            "qwen-local": "fallback: local diffusers QwenImage21Pipeline (default Qwen/Qwen-Image-2.1-Turbo); needs a GPU + `pip install projtex[local]`",
            "hf": "Hugging Face Inference Providers image-to-image (default Qwen/Qwen-Image-Edit-2511); needs HF login/HF_TOKEN",
            "fibo-space": "Bria Fibo 2.0 Edit Turbo via its official HF Space (ZeroGPU quota of your HF login); native masks; CC BY-NC 4.0",
            "files": "writes requests to <project>/exchange and waits for <view>.png (any external generator)",
            "oracle": "re-renders the source mesh's own texture (accuracy tests)",
            "mock": "procedural pattern (smoke tests)"},
                "removed": __import__("projtex.backends", fromlist=["REMOVED"]).REMOVED}
    if c == "schema":
        return _schema(ap)
    if c == "mcp":
        from .mcp_server import main as mcp_main

        mcp_main()
        return None
    if c == "init":
        p = Project.create(args.project, args.mesh, mode=args.mode, uv=args.uv, up=args.up, overwrite=args.overwrite,
                           texture_size=args.texture_size, view_size=args.view_size)
        return {"project": str(p.root), "mesh": p.data["mesh"], "settings": p.settings}
    if c == "run":
        return _run_all(args)
    p = Project.open(args.project)
    if c == "status":
        return p.status()
    if c == "config":
        kw = dict(prompt=args.prompt, subject=args.subject, negative_prompt=args.negative_prompt,
                  lighting=args.lighting, raw_prompt=args.raw_prompt, style_reference=args.style_reference,
                  preset=args.preset, texture_size=args.texture_size, view_size=args.view_size,
                  preview_texture_size=args.preview_texture_size, seed=args.seed, front_azimuth=args.front_azimuth,
                  skymap=p._import_file(args.skymap, "refs") if args.skymap else None)
        if args.fusion:
            kw["fusion"] = _kv(args.fusion)
        kw.update(_kv(args.set))
        b = _backend(args)
        if b:
            kw["backend"] = b
        return p.configure(**kw)
    if c == "plan":
        focus = "default" if args.focus is None else _focus(args.focus)
        return p.plan(preset=args.preset, n_views=args.views, focus=focus, gaps=not args.no_gaps,
                      max_gap_views=args.max_gap_views, fov=args.fov, probes=args.probes, eye_height=args.eye_height,
                      replace=not args.append, levels=args.levels, region=args.region)
    if c == "view":
        if args.remove:
            return p.remove_view(args.name)
        enabled = True if args.enable else (False if args.disable else None)
        return p.set_view(args.name, enabled=enabled, priority=args.priority, reference=args.reference,
                          prompt_extra=args.prompt_extra, seed=args.seed)
    if c == "add-view":
        return p.add_view(args.name, eye=_floats(args.eye), target=_floats(args.target), fov=args.fov,
                          focus=_focus(args.focus), priority=args.priority, index=args.index)
    if c == "generate":
        return p.generate(views=args.views.split(",") if args.views else None, force=args.force,
                          backend=_backend(args), bake=not args.no_bake)
    if c == "fill-gaps":
        return p.fill_gaps(max_views=args.max_views, generate=not args.plan_only, backend=_backend(args),
                           rounds=args.rounds)
    if c == "condition":
        return p.condition_image(args.name)
    if c == "import-view":
        return p.import_view(args.name, args.image)
    if c == "bake":
        return p.bake(resolution=args.resolution)
    if c == "export":
        return p.export(path=args.out, fmt=args.format, resolution=args.resolution, unlit=args.unlit)
    if c == "inspect":
        return p.inspect(max_suggestions=args.max_suggestions)
    if c == "retouch":
        focus, mask, mode = _focus(args.focus), _mask(args.mask), args.mode
        if args.suggestion is not None:
            sug = p.inspect(write=False)["suggestions"]
            if args.suggestion >= len(sug):
                raise SystemExit(f"only {len(sug)} suggestions")
            s = sug[args.suggestion]["retouch"]
            focus = focus or s["focus"]
            mask = mask or s["mask"]
            if "--grow" not in sys.argv:
                args.grow = s.get("grow_px", 0)
            if "--mode" not in sys.argv:
                mode = s["mode"]
        camera = json.loads(args.camera) if args.camera else None
        if args.render_camera:
            camera = json.loads((p.root / "renders" / "cameras.json").read_text())[args.render_camera]
        return p.retouch(prompt=args.prompt, mode=mode, view=args.view, focus=focus,
                         camera=camera, mask=mask,
                         capture_size=args.capture_size, output_size=args.output_size, softness=args.softness,
                         mark=args.mark, backend=_backend(args), seed=args.seed, name=args.name,
                         color_match=not args.no_color_match, strength=args.strength, grow_px=args.grow)
    if c == "pick":
        px = [[float(v) for v in t.split(",")] for t in args.pixels.split(";")]
        camera = json.loads(args.camera) if args.camera else None
        if args.render_camera:
            camera = json.loads((p.root / "renders" / "cameras.json").read_text())[args.render_camera]
        return p.pick(px, view=args.view, camera=camera, size=args.size)
    if c == "undo":
        return p.undo_retouch()
    if c == "layer":
        enabled = True if args.enable else (False if args.disable else None)
        return p.set_layer(args.name, enabled=enabled, strength=args.strength)
    if c == "render":
        out = p.render(out_dir=args.out_dir, turntable=args.turntable, size=args.size, mode=args.mode,
                       elevation=args.elevation, views=args.views.split(",") if args.views else None,
                       compare_source=args.compare_source, video=args.video)
        if args.sheet:
            from .report import views_sheet

            out["views_sheet"] = views_sheet(p)
        return out
    raise SystemExit(f"unknown command {c}")


def _run_all(args) -> dict:
    from .project import Project
    from .report import views_sheet

    p = Project.create(args.project, args.mesh, mode=args.mode, uv=args.uv, up=args.up, overwrite=args.overwrite,
                       texture_size=args.texture_size, view_size=args.view_size)
    p.configure(prompt=args.prompt, subject=args.subject, style_reference=args.style_reference,
                lighting=args.lighting, backend=_backend(args))
    focus = "default" if args.focus is None else _focus(args.focus)
    plan = p.plan(preset=args.preset, probes=args.probes, focus=focus)
    gen = p.generate()
    out = {"project": str(p.root), "plan": {k: v for k, v in plan.items() if k != "views"},
           "views": [r["view"] + ":" + r["status"] for r in gen["results"]]}
    if gen.get("pending"):
        out["pending"] = gen["pending"]
        return out
    retouched = []
    for i in range(args.auto_retouch):
        sug = p.inspect(write=False)["suggestions"]
        if not sug:
            break
        s = sug[0]["retouch"]
        retouched.append(p.retouch(prompt=args.prompt if s["mode"] == "replace" else "", mode=s["mode"],
                                   focus=s["focus"], mask=s["mask"], softness=0.5, grow_px=s.get("grow_px", 4)))
    if retouched:
        out["retouch"] = [r.get("layer") for r in retouched]
    out["bake"] = p.bake()
    out["renders"] = p.render(compare_source=True)
    out["views_sheet"] = views_sheet(p)
    return out


def main(argv: list[str] | None = None) -> int:
    ap = build_parser()
    args = ap.parse_args(argv)
    logging.basicConfig(level=logging.INFO if args.verbose else logging.WARNING, stream=sys.stderr,
                        format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    try:
        res = run(args, ap)
    except SystemExit:
        raise
    except Exception as exc:  # report errors as JSON for agents
        print(json.dumps({"error": f"{type(exc).__name__}: {exc}",
                          "trace": traceback.format_exc(limit=3) if args.verbose else None}))
        return 1
    if res is not None:
        from .project import _json_default

        print(json.dumps(res, indent=2, default=_json_default))
    return 0


if __name__ == "__main__":
    sys.exit(main())