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)