pipeline/pshuman_client.py

"""
pshuman_client.py
=================
Call PSHuman to generate a high-detail 3D face mesh from a portrait image.

Two modes:
- Direct (default when service_url is localhost): runs PSHuman inference.py
  as a subprocess without going through Gradio HTTP.  Avoids the gradio_client
  API-info bug that affects the pshuman Gradio env.
- Remote: uses gradio_client to call a running pshuman_app.py service.

Usage (standalone)
------------------
python -m pipeline.pshuman_client \\
    --image /path/to/portrait.png \\
    --output /tmp/pshuman_face.obj \\
    [--url http://remote-host:7862]   # omit for direct/local mode

Requires: gradio-client (remote mode only)
"""
from __future__ import annotations

import argparse
import glob
import os
import shutil
import subprocess
import time
from pathlib import Path

# Default: assume running on the same instance (local)
_DEFAULT_URL = os.environ.get("PSHUMAN_URL", "http://localhost:7862")

# ── Paths (on the Vast instance) ──────────────────────────────────────────────
PSHUMAN_DIR  = "/root/PSHuman"
CONDA_PYTHON = "/root/miniconda/envs/pshuman/bin/python"
CONFIG       = f"{PSHUMAN_DIR}/configs/inference-768-6view.yaml"
HF_MODEL_DIR = f"{PSHUMAN_DIR}/checkpoints/PSHuman_Unclip_768_6views"
HF_MODEL_HUB = "pengHTYX/PSHuman_Unclip_768_6views"


def _run_pshuman_direct(image_path: str, work_dir: str) -> str:
    """
    Run PSHuman inference.py directly as a subprocess.
    Returns path to the colored OBJ mesh.
    """
    img_dir = os.path.join(work_dir, "input")
    out_dir = os.path.join(work_dir, "out")
    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(out_dir, exist_ok=True)

    scene = "face"
    dst = os.path.join(img_dir, f"{scene}.png")
    shutil.copy(image_path, dst)

    hf_model = HF_MODEL_DIR if Path(HF_MODEL_DIR).exists() else HF_MODEL_HUB

    cmd = [
        CONDA_PYTHON, f"{PSHUMAN_DIR}/inference.py",
        "--config", CONFIG,
        f"pretrained_model_name_or_path={hf_model}",
        f"validation_dataset.root_dir={img_dir}",
        f"save_dir={out_dir}",
        "validation_dataset.crop_size=740",
        "with_smpl=false",
        "num_views=7",
        "save_mode=rgb",
        "seed=42",
    ]

    print(f"[pshuman] Running direct inference: {' '.join(cmd[:4])} ...")
    t0 = time.time()

    # Set CUDA_HOME + extra include dirs so nvdiffrast/torch JIT can compile.
    # On Vast.ai, triposg conda env ships nvcc at bin/nvcc and CUDA headers
    # scattered across site-packages/nvidia/{pkg}/include/ directories.
    env = os.environ.copy()
    if "CUDA_HOME" not in env:
        _triposg   = "/root/miniconda/envs/triposg"
        _targets   = os.path.join(_triposg, "targets", "x86_64-linux")
        _nvcc_bin  = os.path.join(_triposg, "bin")
        _cuda_home = _targets  # has include/cuda_runtime_api.h

        _nvvm_bin  = os.path.join(_triposg, "nvvm", "bin")   # contains cicc
        _nvcc_real = os.path.join(_targets, "bin")            # contains nvcc (real one)

        if (os.path.exists(os.path.join(_cuda_home, "include", "cuda_runtime_api.h"))
                and (os.path.exists(os.path.join(_nvcc_bin, "nvcc"))
                     or os.path.exists(os.path.join(_nvcc_real, "nvcc")))):
            env["CUDA_HOME"] = _cuda_home
            # Build PATH: nvvm/bin (cicc) + targets/.../bin (nvcc real) + conda bin (nvcc wrapper)
            path_parts = []
            if os.path.isdir(_nvvm_bin):
                path_parts.append(_nvvm_bin)
            if os.path.isdir(_nvcc_real):
                path_parts.append(_nvcc_real)
            path_parts.append(_nvcc_bin)
            env["PATH"] = ":".join(path_parts) + ":" + env.get("PATH", "")

            # Collect all nvidia sub-package include dirs (cusparse, cublas, etc.)
            _nvidia_site = os.path.join(_triposg, "lib", "python3.10",
                                        "site-packages", "nvidia")
            _extra_incs = []
            if os.path.isdir(_nvidia_site):
                import glob as _glob
                for _inc in _glob.glob(os.path.join(_nvidia_site, "*/include")):
                    if os.path.isdir(_inc):
                        _extra_incs.append(_inc)
            if _extra_incs:
                _sep = ":"
                _existing = env.get("CPATH", "")
                env["CPATH"] = _sep.join(_extra_incs) + (_sep + _existing if _existing else "")
            print(f"[pshuman] CUDA_HOME={_cuda_home}, {len(_extra_incs)} nvidia include dirs added")

    proc = subprocess.run(
        cmd, cwd=PSHUMAN_DIR,
        capture_output=False,
        text=True,
        timeout=600,
        env=env,
    )
    elapsed = time.time() - t0
    print(f"[pshuman] Inference done in {elapsed:.1f}s (exit={proc.returncode})")

    if proc.returncode != 0:
        raise RuntimeError(f"PSHuman inference failed (exit {proc.returncode})")

    # Locate output OBJ — PSHuman may save relative to its CWD (/root/PSHuman/out/)
    # rather than to the specified save_dir, so check both locations.
    cwd_out_dir = os.path.join(PSHUMAN_DIR, "out", scene)
    patterns = [
        f"{out_dir}/{scene}/result_clr_scale4_{scene}.obj",
        f"{out_dir}/{scene}/result_clr_scale*_{scene}.obj",
        f"{out_dir}/**/*.obj",
        f"{cwd_out_dir}/result_clr_scale*_{scene}.obj",
        f"{cwd_out_dir}/*.obj",
        f"{PSHUMAN_DIR}/out/**/*.obj",
    ]
    obj_path = None
    for pat in patterns:
        hits = sorted(glob.glob(pat, recursive=True))
        if hits:
            colored = [h for h in hits if "clr" in h]
            obj_path = (colored or hits)[-1]
            break

    if not obj_path:
        all_files = list(Path(out_dir).rglob("*"))
        objs = [str(f) for f in all_files if f.suffix in (".obj", ".ply", ".glb")]
        if objs:
            obj_path = objs[-1]
        if not obj_path and Path(cwd_out_dir).exists():
            for f in Path(cwd_out_dir).rglob("*.obj"):
                obj_path = str(f)
                break
        if not obj_path:
            raise FileNotFoundError(
                f"No mesh output found in {out_dir}. "
                f"Files: {[str(f) for f in all_files[:20]]}"
            )

    print(f"[pshuman] Output mesh: {obj_path}")
    return obj_path


def generate_pshuman_mesh(
    image_path: str,
    output_path: str,
    service_url: str = _DEFAULT_URL,
    timeout: float = 600.0,
) -> str:
    """
    Generate a PSHuman face mesh and save it to *output_path*.

    When service_url points to localhost, PSHuman inference.py is run directly
    (no Gradio HTTP, avoids gradio_client API-info bug).
    For remote URLs, gradio_client is used.

    Parameters
    ----------
    image_path   : local PNG/JPG path of the portrait
    output_path  : where to save the downloaded OBJ
    service_url  : base URL of pshuman_app.py, or "direct" to skip HTTP
    timeout      : seconds to wait for inference (used in remote mode)

    Returns
    -------
    output_path  (convenience)
    """
    import tempfile

    output_path = str(output_path)
    os.makedirs(Path(output_path).parent, exist_ok=True)

    is_local = (
        "localhost" in service_url
        or "127.0.0.1" in service_url
        or service_url.strip().lower() == "direct"
        or not service_url.strip()
    )

    if is_local:
        # ── Direct mode: run subprocess ───────────────────────────────────────
        print(f"[pshuman] Direct mode (no HTTP) — running inference on {image_path}")
        work_dir = tempfile.mkdtemp(prefix="pshuman_direct_")
        obj_tmp  = _run_pshuman_direct(image_path, work_dir)
    else:
        # ── Remote mode: call Gradio service ──────────────────────────────────
        try:
            from gradio_client import Client
        except ImportError:
            raise ImportError("pip install gradio-client")

        print(f"[pshuman] Connecting to {service_url}")
        client = Client(service_url)

        print(f"[pshuman] Submitting: {image_path}")
        result = client.predict(
            image=image_path,
            api_name="/gradio_generate_face",
        )

        if isinstance(result, (list, tuple)):
            obj_tmp = result[0]
            status  = result[1] if len(result) > 1 else "ok"
        elif isinstance(result, dict):
            obj_tmp = result.get("obj_path") or result.get("value")
            status  = result.get("status", "ok")
        else:
            obj_tmp = result
            status  = "ok"

        if not obj_tmp or "Error" in str(status):
            raise RuntimeError(f"PSHuman service error: {status}")

        if isinstance(obj_tmp, dict):
            obj_tmp = obj_tmp.get("path") or obj_tmp.get("name") or str(obj_tmp)

        work_dir = str(Path(str(obj_tmp)).parent)

    # ── Copy OBJ + companions to output location ───────────────────────────
    shutil.copy(str(obj_tmp), output_path)
    print(f"[pshuman] Saved OBJ -> {output_path}")

    src_dir = Path(str(obj_tmp)).parent
    out_dir = Path(output_path).parent
    for ext in ("*.mtl", "*.png", "*.jpg"):
        for f in src_dir.glob(ext):
            dest = out_dir / f.name
            if not dest.exists():
                shutil.copy(str(f), str(dest))

    return output_path


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Generate PSHuman face mesh from portrait image"
    )
    parser.add_argument("--image",   required=True, help="Portrait image path")
    parser.add_argument("--output",  required=True, help="Output OBJ path")
    parser.add_argument(
        "--url", default=_DEFAULT_URL,
        help="PSHuman service URL, or 'direct' to run inference locally "
             "(default: http://localhost:7862 → auto-selects direct mode)",
    )
    parser.add_argument("--timeout", type=float, default=600.0)
    args = parser.parse_args()

    generate_pshuman_mesh(
        image_path  = args.image,
        output_path = args.output,
        service_url = args.url,
        timeout     = args.timeout,
    )


if __name__ == "__main__":
    main()