scripts/mesh_generator.py

"""
Mesh generator: text/image → 3D mesh.
- Preferred: Hunyuan3D-2 full pipeline (text → HunyuanDiT → shape → texture). Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2.
- Fallback: TripoSR (image→mesh). Expects TripoSR repo at project_root/TripoSR. Text→mesh uses SD for text→image then TripoSR.
"""

import os
import subprocess
import sys
import time
from pathlib import Path


def find_triposr_root(project_root: str | None = None) -> str | None:
    """Locate TripoSR repo: ./TripoSR or ../TripoSR from script dir."""
    if project_root is None:
        project_root = str(Path(__file__).resolve().parent.parent)
    candidates = [
        os.path.join(project_root, "TripoSR"),
        os.path.join(project_root, "..", "TripoSR"),
    ]
    for p in candidates:
        run_py = os.path.join(p, "run.py")
        if os.path.isfile(run_py):
            return p
    return None


def _obj_texture_to_glb(obj_path: str, texture_path: str, glb_path: str) -> None:
    """
    Convert OBJ + texture.png to a valid GLB with embedded texture.
    TripoSR's xatlas.export writes OBJ format regardless of extension; use this to produce real GLB.
    """
    import trimesh
    from PIL import Image
    mesh = trimesh.load(obj_path, file_type="obj", process=False)
    if not isinstance(mesh, trimesh.Trimesh):
        mesh = mesh.dump(concatenate=True) if hasattr(mesh, "dump") else None
    if mesh is None:
        return
    # Ensure texture image is set (xatlas may not write MTL)
    if hasattr(mesh, "visual") and mesh.visual is not None and getattr(mesh.visual, "uv", None) is not None:
        try:
            img = Image.open(texture_path)
            from trimesh.visual import TextureVisuals
            mesh.visual = TextureVisuals(uv=mesh.visual.uv, image=img)
        except Exception:
            pass
    mesh.export(glb_path)


def _is_obj_content(path: str) -> bool:
    """Return True if file content looks like OBJ (TripoSR xatlas writes OBJ even when path is .glb)."""
    try:
        with open(path, "rb") as f:
            return f.read(2).strip() == b"v"
    except Exception:
        return False


def _smooth_mesh_file(mesh_path: str, iterations: int = 3, lamb: float = 0.5) -> None:
    """Apply light Laplacian smoothing to a mesh file in-place. Preserves UVs."""
    import trimesh
    loaded = trimesh.load(mesh_path, force="mesh", process=False)
    if isinstance(loaded, trimesh.Trimesh):
        mesh = loaded
    elif hasattr(loaded, "dump"):
        mesh = loaded.dump(concatenate=True)
    else:
        return
    if mesh is None:
        return
    try:
        trimesh.smoothing.filter_laplacian(mesh, lamb=lamb, iterations=iterations)
    except Exception:
        return
    mesh.export(mesh_path)


def generate_mesh_from_image(
    image_path: str,
    output_dir: str = "outputs",
    mesh_format: str = "glb",
    triposr_root: str | None = None,
    device: str | None = None,
    use_hunyuan3d2: bool | None = None,
    mc_resolution: int = 512,
    bake_texture: bool = True,
    texture_resolution: int = 2048,
    smooth_mesh: bool = True,
) -> tuple[str | None, float, str]:
    """
    Image → 3D mesh. Uses Hunyuan3D-2 when available (use_hunyuan3d2=True or repo found), else TripoSR.
    mc_resolution: marching cubes grid (256=faster, 512=higher quality). bake_texture: use texture atlas.
    smooth_mesh: apply light Laplacian smoothing. Returns (path_to_mesh, inference_time_sec, message).
    """
    project_root = str(Path(__file__).resolve().parent.parent)
    hunyuan_root = find_hunyuan3d2_root(project_root)
    if use_hunyuan3d2 is True or (use_hunyuan3d2 is None and hunyuan_root is not None):
        if hunyuan_root:
            from scripts.hunyuan3d_text_to_mesh import generate_mesh_from_image_hunyuan3d2
            return generate_mesh_from_image_hunyuan3d2(
                image_path,
                output_dir=output_dir,
                mesh_format=mesh_format,
                seed=42,
                with_texture=True,
                hunyuan_root=hunyuan_root,
            )
        if use_hunyuan3d2 is True:
            return (None, 0.0, "Hunyuan3D-2 repo not found. Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2 (see README).")

    triposr_root = triposr_root or find_triposr_root(project_root)
    if not triposr_root:
        return (
            None,
            0.0,
            "TripoSR not found. Clone it: git clone https://github.com/VAST-AI-Research/TripoSR.git",
        )

    device = device or _infer_device()
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    # When baking texture, xatlas.export() always writes OBJ format (ignores .glb extension).
    # Ask for OBJ when bake_texture + glb, then we convert OBJ+texture to real GLB.
    triposr_format = "obj" if (bake_texture and mesh_format == "glb") else mesh_format
    run_py = os.path.join(triposr_root, "run.py")
    cmd = [
        sys.executable,
        run_py,
        image_path,
        "--output-dir",
        output_dir,
        "--model-save-format",
        triposr_format,
        "--device",
        device,
        "--mc-resolution",
        str(mc_resolution),
    ]
    if bake_texture:
        cmd += ["--bake-texture", "--texture-resolution", str(texture_resolution)]
    # Use only the current env (venv) so torch and torchvision match; avoids
    # "operator torchvision::nms does not exist" when user site-packages mixed in.
    env = os.environ.copy()
    env["PYTHONNOUSERSITE"] = "1"
    # TripoSR downloads its model from Hugging Face Hub on first run; ensure token is available in subprocess
    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if hf_token:
        env["HF_TOKEN"] = hf_token
        env["HUGGING_FACE_HUB_TOKEN"] = hf_token

    t0 = time.perf_counter()
    try:
        result = subprocess.run(
            cmd,
            cwd=triposr_root,
            env=env,
            capture_output=True,
            text=True,
            timeout=600,
        )
        t1 = time.perf_counter()
        if result.returncode != 0:
            err_text = (result.stderr or result.stdout or "unknown").strip()
            if "403" in err_text or "forbidden" in err_text.lower():
                return (
                    None,
                    t1 - t0,
                    "TripoSR got 403 from Hugging Face Hub (model download). "
                    "Set HF_TOKEN in this Space: Settings → Variables and secrets (get a token at huggingface.co/settings/tokens).",
                )
            return (None, t1 - t0, f"TripoSR failed: {err_text}")
        # Output is output_dir/0/mesh.glb or mesh.obj
        out_subdir = os.path.join(output_dir, "0")
        mesh_path = os.path.join(out_subdir, f"mesh.{triposr_format}")
        if not os.path.isfile(mesh_path):
            return (None, t1 - t0, f"TripoSR did not produce {mesh_path}")
        # If we asked for GLB but TripoSR wrote OBJ (bake_texture), convert to real GLB
        if bake_texture and mesh_format == "glb":
            texture_path = os.path.join(out_subdir, "texture.png")
            glb_path = os.path.join(out_subdir, "mesh.glb")
            try:
                _obj_texture_to_glb(mesh_path, texture_path, glb_path)
                mesh_path = glb_path
            except Exception:
                pass  # keep mesh.obj path if conversion fails
        # If existing file is misnamed (OBJ content in .glb), repair it
        elif mesh_format == "glb" and _is_obj_content(mesh_path):
            texture_path = os.path.join(out_subdir, "texture.png")
            try:
                _obj_texture_to_glb(mesh_path, texture_path, mesh_path)
            except Exception:
                pass
        if smooth_mesh:
            try:
                _smooth_mesh_file(mesh_path, iterations=3, lamb=0.5)
            except Exception:
                pass
        return (os.path.abspath(mesh_path), t1 - t0, "OK")
    except subprocess.TimeoutExpired:
        t1 = time.perf_counter()
        return (None, t1 - t0, "TripoSR timed out (10 min). Try lower resolution (256) or disable bake texture.")
    except Exception as e:
        t1 = time.perf_counter()
        return (None, t1 - t0, str(e))


def find_hunyuan3d2_root(project_root: str | None = None) -> str | None:
    """Locate Hunyuan3D-2 repo. Delegates to hunyuan3d_text_to_mesh."""
    try:
        from scripts.hunyuan3d_text_to_mesh import find_hunyuan3d2_root as _find
        return _find(project_root)
    except Exception:
        return None


def _infer_device() -> str:
    """Use GPU if available, else CPU (for CPU-only Spaces)."""
    try:
        import torch
        return "cuda:0" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


def generate_mesh_from_text(
    prompt: str,
    output_dir: str = "outputs",
    mesh_format: str = "glb",
    seed: int | None = None,
    use_hunyuan3d2: bool | None = None,
    mc_resolution: int = 512,
    bake_texture: bool = True,
    texture_resolution: int = 2048,
    smooth_mesh: bool = True,
    device: str | None = None,
) -> tuple[str | None, float, str]:
    """
    Text → 3D mesh. Uses Hunyuan3D-2 full pipeline when available (use_hunyuan3d2=True or repo found),
    else TripoSR (SD for text→image then TripoSR). Returns (path_to_mesh, total_time_sec, message).
    """
    _root = str(Path(__file__).resolve().parent.parent)
    if _root not in sys.path:
        sys.path.insert(0, _root)
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    seed = seed if seed is not None else 42

    # Prefer Hunyuan3D-2 when requested or when it's the only backend available
    hunyuan_root = find_hunyuan3d2_root(_root)
    if use_hunyuan3d2 is True or (use_hunyuan3d2 is None and hunyuan_root is not None):
        if hunyuan_root:
            from scripts.hunyuan3d_text_to_mesh import generate_mesh_from_text_hunyuan3d2
            return generate_mesh_from_text_hunyuan3d2(
                prompt,
                output_dir=output_dir,
                mesh_format=mesh_format,
                seed=seed,
                with_texture=True,
                hunyuan_root=hunyuan_root,
            )
        if use_hunyuan3d2 is True:
            return (None, 0.0, "Hunyuan3D-2 repo not found. Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2 (see README).")

    # TripoSR path: text → image (SD) → mesh
    from scripts.text_to_image import text_to_image
    t0 = time.perf_counter()
    try:
        image_path, _ = text_to_image(prompt, output_dir=output_dir, seed=seed)
    except Exception as e:
        return (None, 0.0, f"Text-to-image failed: {e}")
    device = device or _infer_device()
    mesh_path, mesh_time, msg = generate_mesh_from_image(
        image_path,
        output_dir=os.path.join(output_dir, "mesh_run"),
        mesh_format=mesh_format,
        mc_resolution=mc_resolution,
        bake_texture=bake_texture,
        texture_resolution=texture_resolution,
        smooth_mesh=smooth_mesh,
        device=device,
    )
    total_time = time.perf_counter() - t0
    if mesh_path:
        return (mesh_path, total_time, msg)
    return (None, total_time, msg)