app.py

import logging
import os
from pathlib import Path
import site
from typing import Any, Tuple
from uuid import uuid4

import gradio as gr
from PIL import Image
import trimesh

try:
    import spaces
except ImportError:
    class _FakeGPU:
        def __init__(self, duration: int = 60):
            self.duration = duration

        def __call__(self, fn):
            return fn

    class spaces:  # type: ignore
        GPU = _FakeGPU


logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger("gamemaster_attire_transfer")

MODEL_ID = os.environ.get("HY3D_TEXGEN_MODEL", "tencent/Hunyuan3D-2")
OUTPUT_DIR = Path(os.environ.get("GM_ATTIRE_OUTPUT_DIR", "/tmp/gamemaster_attire_transfer"))
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
MODEL_CACHE_ROOT = Path(os.environ.get("HY3DGEN_MODELS", "/tmp/hy3dgen_models"))
MODEL_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
os.environ["HY3DGEN_MODELS"] = MODEL_CACHE_ROOT.as_posix()

_TEXGEN_PIPELINE: Any = None
_BG_REMOVER = None


def _prefetch_hunyuan_models() -> None:
    from huggingface_hub import snapshot_download

    local_dir = MODEL_CACHE_ROOT / MODEL_ID
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
    LOGGER.info("Prefetching Hunyuan model files into %s", local_dir)
    snapshot_download(
        repo_id=MODEL_ID,
        local_dir=local_dir.as_posix(),
        allow_patterns=[
            "hunyuan3d-delight-v2-0/*",
            "hunyuan3d-paint-v2-0/*",
        ],
        token=token,
    )


def _prepend_env_path(var_name: str, entries: list[str], separator: str) -> None:
    cleaned = [entry for entry in entries if entry]
    if not cleaned:
        return
    current = os.environ.get(var_name, "")
    merged = separator.join(cleaned + ([current] if current else []))
    os.environ[var_name] = merged


def _configure_cuda_runtime_paths() -> None:
    # On HF containers, CUDA runtime and nvcc are often installed under site-packages/nvidia/*.
    # We expose these folders so extension build/import can find nvcc and shared libs.
    lib_dirs: list[str] = []
    bin_dirs: list[str] = []
    cuda_home: str | None = None

    for base in site.getsitepackages():
        nvidia_root = Path(base) / "nvidia"
        if nvidia_root.exists():
            for lib_dir in nvidia_root.glob("*/lib"):
                if lib_dir.is_dir():
                    lib_dirs.append(str(lib_dir))
            for bin_dir in nvidia_root.glob("*/bin"):
                if bin_dir.is_dir():
                    bin_dirs.append(str(bin_dir))
            nvcc_root = nvidia_root / "cuda_nvcc"
            if nvcc_root.exists() and (nvcc_root / "bin").exists():
                cuda_home = str(nvcc_root)
        torch_lib = Path(base) / "torch" / "lib"
        if torch_lib.exists():
            lib_dirs.append(str(torch_lib))

    _prepend_env_path("LD_LIBRARY_PATH", lib_dirs, ":")
    _prepend_env_path("PATH", bin_dirs, os.pathsep)

    if cuda_home:
        os.environ["CUDA_HOME"] = cuda_home
        os.environ["CUDA_PATH"] = cuda_home


def _ensure_custom_rasterizer_installed() -> None:
    _configure_cuda_runtime_paths()
    try:
        import torch  # noqa: F401
        import custom_rasterizer  # noqa: F401
        return
    except Exception as exc:
        raise RuntimeError(
            "custom_rasterizer failed to import. Ensure the wheel is compatible with the installed torch/CUDA stack."
        ) from exc


_prefetch_hunyuan_models()


def _get_pipeline():
    global _TEXGEN_PIPELINE
    if _TEXGEN_PIPELINE is None:
        _ensure_custom_rasterizer_installed()
        from hy3dgen.texgen import Hunyuan3DPaintPipeline
        LOGGER.info("Loading texgen pipeline from %s", MODEL_ID)
        _TEXGEN_PIPELINE = Hunyuan3DPaintPipeline.from_pretrained(MODEL_ID)
        try:
            _TEXGEN_PIPELINE.enable_model_cpu_offload()
            LOGGER.info("Enabled model CPU offload for texgen pipeline")
        except Exception as exc:
            LOGGER.warning("Could not enable CPU offload: %s", exc)
    return _TEXGEN_PIPELINE


def _get_background_remover():
    global _BG_REMOVER
    if _BG_REMOVER is None:
        from hy3dgen.rembg import BackgroundRemover
        _BG_REMOVER = BackgroundRemover()
    return _BG_REMOVER


def _load_single_mesh(mesh_path: str) -> trimesh.Trimesh:
    loaded = trimesh.load(mesh_path, force="scene")
    if isinstance(loaded, trimesh.Scene):
        if not loaded.geometry:
            raise ValueError("The uploaded GLB contains no geometry.")
        merged = loaded.to_geometry()
        if not isinstance(merged, trimesh.Trimesh):
            raise ValueError("Failed to convert GLB scene to a single mesh.")
        return merged
    if isinstance(loaded, trimesh.Trimesh):
        return loaded
    raise ValueError("Unsupported mesh type. Please upload a valid GLB mesh.")


def _prepare_reference_image(image_path: str, remove_background: bool) -> Image.Image:
    image = Image.open(image_path).convert("RGBA")
    if remove_background:
        image = _get_background_remover()(image)
    return image


@spaces.GPU(duration=60)
def transfer_attire(
    reference_image_path: str,
    source_mesh_path: str,
    remove_background: bool,
) -> Tuple[str, str]:
    try:
        if not reference_image_path:
            raise gr.Error("Please upload a reference image.")
        if not source_mesh_path:
            raise gr.Error("Please upload a source GLB mesh.")

        import torch
        LOGGER.info("CUDA available=%s, device_count=%s", torch.cuda.is_available(), torch.cuda.device_count())

        image = _prepare_reference_image(reference_image_path, remove_background=remove_background)
        mesh = _load_single_mesh(source_mesh_path)

        textured_mesh = _get_pipeline()(mesh, image)

        out_dir = OUTPUT_DIR / str(uuid4())
        out_dir.mkdir(parents=True, exist_ok=True)
        out_glb = out_dir / "textured_mesh.glb"
        textured_mesh.export(out_glb.as_posix(), include_normals=True)

        return out_glb.as_posix(), out_glb.as_posix()
    except Exception as exc:
        LOGGER.exception("transfer_attire failed")
        raise gr.Error(f"transfer_attire failed: {exc}") from exc


with gr.Blocks(title="GameMaster Attire Transfer (Hunyuan3D-2)") as demo:
    gr.Markdown(
        """
        # GameMaster Attire Transfer (Hunyuan3D-2)
        Upload a reference image and a GLB mesh. The app transfers the clothing/material style from
        the image onto the uploaded mesh and exports a textured GLB.
        """
    )

    with gr.Row():
        with gr.Column():
            reference_image = gr.Image(type="filepath", label="Reference Image")
            source_mesh = gr.File(file_types=[".glb"], type="filepath", label="Source Mesh (GLB)")
            remove_bg = gr.Checkbox(value=True, label="Remove reference image background")
            run_button = gr.Button("Transfer Attire", variant="primary")
        with gr.Column():
            output_preview = gr.Model3D(label="Textured Output Preview")
            output_glb = gr.File(label="Download Textured GLB")

    run_button.click(
        transfer_attire,
        inputs=[reference_image, source_mesh, remove_bg],
        outputs=[output_preview, output_glb],
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)