evoneural/evoneuralIn3D-app

1"""2Mesh generator: text/image → 3D mesh.3- Preferred: Hunyuan3D-2 full pipeline (text → HunyuanDiT → shape → texture). Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2.4- Fallback: TripoSR (image→mesh). Expects TripoSR repo at project_root/TripoSR. Text→mesh uses SD for text→image then TripoSR.5"""6 7import os8import subprocess9import sys10import time11from pathlib import Path12 13 14def find_triposr_root(project_root: str | None = None) -> str | None:15    """Locate TripoSR repo: ./TripoSR or ../TripoSR from script dir."""16    if project_root is None:17        project_root = str(Path(__file__).resolve().parent.parent)18    candidates = [19        os.path.join(project_root, "TripoSR"),20        os.path.join(project_root, "..", "TripoSR"),21    ]22    for p in candidates:23        run_py = os.path.join(p, "run.py")24        if os.path.isfile(run_py):25            return p26    return None27 28 29def _obj_texture_to_glb(obj_path: str, texture_path: str, glb_path: str) -> None:30    """31    Convert OBJ + texture.png to a valid GLB with embedded texture.32    TripoSR's xatlas.export writes OBJ format regardless of extension; use this to produce real GLB.33    """34    import trimesh35    from PIL import Image36    mesh = trimesh.load(obj_path, file_type="obj", process=False)37    if not isinstance(mesh, trimesh.Trimesh):38        mesh = mesh.dump(concatenate=True) if hasattr(mesh, "dump") else None39    if mesh is None:40        return41    # Ensure texture image is set (xatlas may not write MTL)42    if hasattr(mesh, "visual") and mesh.visual is not None and getattr(mesh.visual, "uv", None) is not None:43        try:44            img = Image.open(texture_path)45            from trimesh.visual import TextureVisuals46            mesh.visual = TextureVisuals(uv=mesh.visual.uv, image=img)47        except Exception:48            pass49    mesh.export(glb_path)50 51 52def _is_obj_content(path: str) -> bool:53    """Return True if file content looks like OBJ (TripoSR xatlas writes OBJ even when path is .glb)."""54    try:55        with open(path, "rb") as f:56            return f.read(2).strip() == b"v"57    except Exception:58        return False59 60 61def _smooth_mesh_file(mesh_path: str, iterations: int = 3, lamb: float = 0.5) -> None:62    """Apply light Laplacian smoothing to a mesh file in-place. Preserves UVs."""63    import trimesh64    loaded = trimesh.load(mesh_path, force="mesh", process=False)65    if isinstance(loaded, trimesh.Trimesh):66        mesh = loaded67    elif hasattr(loaded, "dump"):68        mesh = loaded.dump(concatenate=True)69    else:70        return71    if mesh is None:72        return73    try:74        trimesh.smoothing.filter_laplacian(mesh, lamb=lamb, iterations=iterations)75    except Exception:76        return77    mesh.export(mesh_path)78 79 80def generate_mesh_from_image(81    image_path: str,82    output_dir: str = "outputs",83    mesh_format: str = "glb",84    triposr_root: str | None = None,85    device: str | None = None,86    use_hunyuan3d2: bool | None = None,87    mc_resolution: int = 512,88    bake_texture: bool = True,89    texture_resolution: int = 2048,90    smooth_mesh: bool = True,91) -> tuple[str | None, float, str]:92    """93    Image → 3D mesh. Uses Hunyuan3D-2 when available (use_hunyuan3d2=True or repo found), else TripoSR.94    mc_resolution: marching cubes grid (256=faster, 512=higher quality). bake_texture: use texture atlas.95    smooth_mesh: apply light Laplacian smoothing. Returns (path_to_mesh, inference_time_sec, message).96    """97    project_root = str(Path(__file__).resolve().parent.parent)98    hunyuan_root = find_hunyuan3d2_root(project_root)99    if use_hunyuan3d2 is True or (use_hunyuan3d2 is None and hunyuan_root is not None):100        if hunyuan_root:101            from scripts.hunyuan3d_text_to_mesh import generate_mesh_from_image_hunyuan3d2102            return generate_mesh_from_image_hunyuan3d2(103                image_path,104                output_dir=output_dir,105                mesh_format=mesh_format,106                seed=42,107                with_texture=True,108                hunyuan_root=hunyuan_root,109            )110        if use_hunyuan3d2 is True:111            return (None, 0.0, "Hunyuan3D-2 repo not found. Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2 (see README).")112 113    triposr_root = triposr_root or find_triposr_root(project_root)114    if not triposr_root:115        return (116            None,117            0.0,118            "TripoSR not found. Clone it: git clone https://github.com/VAST-AI-Research/TripoSR.git",119        )120 121    device = device or _infer_device()122    Path(output_dir).mkdir(parents=True, exist_ok=True)123    # When baking texture, xatlas.export() always writes OBJ format (ignores .glb extension).124    # Ask for OBJ when bake_texture + glb, then we convert OBJ+texture to real GLB.125    triposr_format = "obj" if (bake_texture and mesh_format == "glb") else mesh_format126    run_py = os.path.join(triposr_root, "run.py")127    cmd = [128        sys.executable,129        run_py,130        image_path,131        "--output-dir",132        output_dir,133        "--model-save-format",134        triposr_format,135        "--device",136        device,137        "--mc-resolution",138        str(mc_resolution),139    ]140    if bake_texture:141        cmd += ["--bake-texture", "--texture-resolution", str(texture_resolution)]142    # Use only the current env (venv) so torch and torchvision match; avoids143    # "operator torchvision::nms does not exist" when user site-packages mixed in.144    env = os.environ.copy()145    env["PYTHONNOUSERSITE"] = "1"146    # TripoSR downloads its model from Hugging Face Hub on first run; ensure token is available in subprocess147    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")148    if hf_token:149        env["HF_TOKEN"] = hf_token150        env["HUGGING_FACE_HUB_TOKEN"] = hf_token151 152    t0 = time.perf_counter()153    try:154        result = subprocess.run(155            cmd,156            cwd=triposr_root,157            env=env,158            capture_output=True,159            text=True,160            timeout=600,161        )162        t1 = time.perf_counter()163        if result.returncode != 0:164            err_text = (result.stderr or result.stdout or "unknown").strip()165            if "403" in err_text or "forbidden" in err_text.lower():166                return (167                    None,168                    t1 - t0,169                    "TripoSR got 403 from Hugging Face Hub (model download). "170                    "Set HF_TOKEN in this Space: Settings → Variables and secrets (get a token at huggingface.co/settings/tokens).",171                )172            return (None, t1 - t0, f"TripoSR failed: {err_text}")173        # Output is output_dir/0/mesh.glb or mesh.obj174        out_subdir = os.path.join(output_dir, "0")175        mesh_path = os.path.join(out_subdir, f"mesh.{triposr_format}")176        if not os.path.isfile(mesh_path):177            return (None, t1 - t0, f"TripoSR did not produce {mesh_path}")178        # If we asked for GLB but TripoSR wrote OBJ (bake_texture), convert to real GLB179        if bake_texture and mesh_format == "glb":180            texture_path = os.path.join(out_subdir, "texture.png")181            glb_path = os.path.join(out_subdir, "mesh.glb")182            try:183                _obj_texture_to_glb(mesh_path, texture_path, glb_path)184                mesh_path = glb_path185            except Exception:186                pass  # keep mesh.obj path if conversion fails187        # If existing file is misnamed (OBJ content in .glb), repair it188        elif mesh_format == "glb" and _is_obj_content(mesh_path):189            texture_path = os.path.join(out_subdir, "texture.png")190            try:191                _obj_texture_to_glb(mesh_path, texture_path, mesh_path)192            except Exception:193                pass194        if smooth_mesh:195            try:196                _smooth_mesh_file(mesh_path, iterations=3, lamb=0.5)197            except Exception:198                pass199        return (os.path.abspath(mesh_path), t1 - t0, "OK")200    except subprocess.TimeoutExpired:201        t1 = time.perf_counter()202        return (None, t1 - t0, "TripoSR timed out (10 min). Try lower resolution (256) or disable bake texture.")203    except Exception as e:204        t1 = time.perf_counter()205        return (None, t1 - t0, str(e))206 207 208def find_hunyuan3d2_root(project_root: str | None = None) -> str | None:209    """Locate Hunyuan3D-2 repo. Delegates to hunyuan3d_text_to_mesh."""210    try:211        from scripts.hunyuan3d_text_to_mesh import find_hunyuan3d2_root as _find212        return _find(project_root)213    except Exception:214        return None215 216 217def _infer_device() -> str:218    """Use GPU if available, else CPU (for CPU-only Spaces)."""219    try:220        import torch221        return "cuda:0" if torch.cuda.is_available() else "cpu"222    except Exception:223        return "cpu"224 225 226def generate_mesh_from_text(227    prompt: str,228    output_dir: str = "outputs",229    mesh_format: str = "glb",230    seed: int | None = None,231    use_hunyuan3d2: bool | None = None,232    mc_resolution: int = 512,233    bake_texture: bool = True,234    texture_resolution: int = 2048,235    smooth_mesh: bool = True,236    device: str | None = None,237) -> tuple[str | None, float, str]:238    """239    Text → 3D mesh. Uses Hunyuan3D-2 full pipeline when available (use_hunyuan3d2=True or repo found),240    else TripoSR (SD for text→image then TripoSR). Returns (path_to_mesh, total_time_sec, message).241    """242    _root = str(Path(__file__).resolve().parent.parent)243    if _root not in sys.path:244        sys.path.insert(0, _root)245    Path(output_dir).mkdir(parents=True, exist_ok=True)246    seed = seed if seed is not None else 42247 248    # Prefer Hunyuan3D-2 when requested or when it's the only backend available249    hunyuan_root = find_hunyuan3d2_root(_root)250    if use_hunyuan3d2 is True or (use_hunyuan3d2 is None and hunyuan_root is not None):251        if hunyuan_root:252            from scripts.hunyuan3d_text_to_mesh import generate_mesh_from_text_hunyuan3d2253            return generate_mesh_from_text_hunyuan3d2(254                prompt,255                output_dir=output_dir,256                mesh_format=mesh_format,257                seed=seed,258                with_texture=True,259                hunyuan_root=hunyuan_root,260            )261        if use_hunyuan3d2 is True:262            return (None, 0.0, "Hunyuan3D-2 repo not found. Set HUNYUAN3D2_ROOT or clone ./Hunyuan3D-2 (see README).")263 264    # TripoSR path: text → image (SD) → mesh265    from scripts.text_to_image import text_to_image266    t0 = time.perf_counter()267    try:268        image_path, _ = text_to_image(prompt, output_dir=output_dir, seed=seed)269    except Exception as e:270        return (None, 0.0, f"Text-to-image failed: {e}")271    device = device or _infer_device()272    mesh_path, mesh_time, msg = generate_mesh_from_image(273        image_path,274        output_dir=os.path.join(output_dir, "mesh_run"),275        mesh_format=mesh_format,276        mc_resolution=mc_resolution,277        bake_texture=bake_texture,278        texture_resolution=texture_resolution,279        smooth_mesh=smooth_mesh,280        device=device,281    )282    total_time = time.perf_counter() - t0283    if mesh_path:284        return (mesh_path, total_time, msg)285    return (None, total_time, msg)286