import os # Force attention backends compatible with the new ZeroGPU (Blackwell) stack. # Must be set BEFORE any trellis import. os.environ.setdefault('ATTN_BACKEND', 'xformers') os.environ.setdefault('SPARSE_ATTN_BACKEND', 'xformers') os.environ.setdefault('SPCONV_ALGO', 'native') import sys import subprocess import tempfile import ctypes import spaces import torch import gradio as gr import shutil from typing import * import numpy as np import imageio from PIL import Image # Build nvdiffrast and diff_gaussian_rasterization from source on first GPU call. CUDA_HOME = "/cuda-image/usr/local/cuda-13.0" CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64") _NVDIFFRAST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "extensions", "nvdiffrast") @spaces.GPU(duration=600) def _first_gpu_setup(): need = {} for name, modname in [ ("nvdiffrast", "nvdiffrast"), ("diff_gaussian_rasterization", "diff_gaussian_rasterization"), ]: try: __import__(modname) except ImportError: need[name] = True if not need: return patch_dir = tempfile.mkdtemp(prefix="torch_cuda_patch_") with open(os.path.join(patch_dir, "sitecustomize.py"), "w") as f: f.write( "try:\n" " import torch.utils.cpp_extension as _c\n" " _c._check_cuda_version = lambda *a, **k: None\n" "except Exception:\n" " pass\n" ) env = os.environ.copy() env["CUDA_HOME"] = CUDA_HOME env["CUDA_PATH"] = CUDA_HOME env["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + env.get("PATH", "") env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "") env["TORCH_CUDA_ARCH_LIST"] = "12.0" # Blackwell sm_120 subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-deps", "setuptools", "wheel", "ninja", "packaging"], ) if "nvdiffrast" in need: subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-build-isolation", "--no-deps", _NVDIFFRAST_DIR], env=env, ) if "diff_gaussian_rasterization" in need: # Hi3DGen actually uses the mip-splatting submodule fork; not the # original graphdeco-inria release on PyPI. mip = tempfile.mkdtemp(prefix="mip_") subprocess.check_call( ["git", "clone", "--recursive", "--depth=1", "https://github.com/autonomousvision/mip-splatting.git", mip], ) subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-build-isolation", "--no-deps", os.path.join(mip, "submodules", "diff-gaussian-rasterization")], env=env, ) _first_gpu_setup() try: ctypes.CDLL(os.path.join(CUDA_LIBDIR, "libcudart.so.13"), mode=ctypes.RTLD_GLOBAL) os.environ["LD_LIBRARY_PATH"] = CUDA_LIBDIR + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "") except OSError: pass from trellis.pipelines import TrellisImageTo3DPipeline from trellis.utils import render_utils import trimesh import tempfile MAX_SEED = np.iinfo(np.int32).max TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') os.makedirs(TMP_DIR, exist_ok=True) from transformers.utils.hub import move_cache move_cache() def preprocess_mesh(mesh_prompt): print("Processing mesh") trimesh_mesh = trimesh.load_mesh(mesh_prompt) trimesh_mesh.export(mesh_prompt+'.glb') return mesh_prompt+'.glb' def preprocess_image(image): if image is None: return None image = pipeline.preprocess_image(image, resolution=1024) return image @spaces.GPU def generate_3d(image, seed=-1, ss_guidance_strength=3, ss_sampling_steps=50, slat_guidance_strength=3, slat_sampling_steps=6,): if image is None: return None, None, None if seed == -1: seed = np.random.randint(0, MAX_SEED) image = pipeline.preprocess_image(image, resolution=1024) normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object') outputs = pipeline.run( normal_image, seed=seed, formats=["mesh",], preprocess_image=False, sparse_structure_sampler_params={ "steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength, }, ) generated_mesh = outputs['mesh'][0] # Save outputs import datetime output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S") os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True) mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb" render_results = render_utils.render_video(generated_mesh, resolution=1024, ssaa=1, num_frames=8, pitch=0.25, inverse_direction=True) def combine_diagonal(color_np, normal_np): # Convert images to numpy arrays h, w, c = color_np.shape # Create a boolean mask that is True for pixels where x > y (diagonally) mask = np.fromfunction(lambda y, x: x > y, (h, w)) mask = mask.astype(bool) mask = np.stack([mask] * c, axis=-1) # Where mask is True take color, else normal combined_np = np.where(mask, color_np, normal_np) return Image.fromarray(combined_np) preview_images = [combine_diagonal(c, n) for c, n in zip(render_results['color'], render_results['normal'])] # Export mesh trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True) trimesh_mesh.export(mesh_path) return preview_images, normal_image, mesh_path, mesh_path def convert_mesh(mesh_path, export_format): """Download the mesh in the selected format.""" if not mesh_path: return None # Create a temporary file to store the mesh data temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False) temp_file_path = temp_file.name new_mesh_path = mesh_path.replace(".glb", f".{export_format}") mesh = trimesh.load_mesh(mesh_path) mesh.export(temp_file_path) # Export to the temporary file return temp_file_path # Return the path to the temporary file # Create the Gradio interface with improved layout with gr.Blocks(css="footer {visibility: hidden}") as demo: gr.Markdown( """

Hi3DGen: High-fidelity 3D Geometry Generation from Images via Normal Bridging

V0.1, Introduced By GAP Lab from CUHKSZ and Game-AIGC Team from ByteDance

""" ) with gr.Row(): gr.Markdown("""

badge-github-stars social

""") with gr.Row(): with gr.Column(scale=1): with gr.Tabs(): with gr.Tab("Single Image"): with gr.Row(): image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil") normal_output = gr.Image(label="Normal Bridge", image_mode="RGBA", type="pil") with gr.Tab("Multiple Images"): gr.Markdown("
Multiple Images functionality is coming soon!
") with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider(-1, MAX_SEED, label="Seed", value=0, step=1) gr.Markdown("#### Stage 1: Sparse Structure Generation") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3, step=0.1) ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=50, step=1) gr.Markdown("#### Stage 2: Structured Latent Generation") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=6, step=1) with gr.Group(): with gr.Row(): gen_shape_btn = gr.Button("Generate Shape", size="lg", variant="primary") # Right column - Output with gr.Column(scale=1): with gr.Tabs(): with gr.Tab("Preview"): output_gallery = gr.Gallery(label="Examples", columns=4, rows=2, object_fit="contain", height="auto",show_label=False) with gr.Tab("3D Model"): with gr.Column(): model_output = gr.Model3D(label="3D Model Preview (Each model is approximately 40MB, may take around 1 minute to load)") with gr.Column(): export_format = gr.Dropdown( choices=["obj", "glb", "ply", "stl"], value="glb", label="File Format" ) download_btn = gr.DownloadButton(label="Export Mesh", interactive=False) image_prompt.upload( preprocess_image, inputs=[image_prompt], outputs=[image_prompt] ) gen_shape_btn.click( generate_3d, inputs=[ image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps ], outputs=[output_gallery, normal_output, model_output, download_btn] ).then( lambda: gr.Button(interactive=True), outputs=[download_btn], ) def update_download_button(mesh_path, export_format): if not mesh_path: return gr.File.update(value=None, interactive=False) download_path = convert_mesh(mesh_path, export_format) return download_path export_format.change( update_download_button, inputs=[model_output, export_format], outputs=[download_btn] ).then( lambda: gr.Button(interactive=True), outputs=[download_btn], ) examples = gr.Examples( examples=[ f'assets/example_image/{image}' for image in os.listdir("assets/example_image") ], inputs=image_prompt, ) gr.Markdown( """ **Acknowledgments**: Hi3DGen is built on the shoulders of giants. We would like to express our gratitude to the open-source research community and the developers of these pioneering projects: - **3D Modeling:** Our 3D Model is finetuned from the SOTA open-source 3D foundation model [Trellis](https://github.com/microsoft/TRELLIS) and we draw inspiration from the teams behind [Rodin](https://hyperhuman.deemos.com/rodin), [Tripo](https://www.tripo3d.ai/app/home), and [Dora](https://github.com/Seed3D/Dora). - **Normal Estimation:** Our Normal Estimation Model builds on the leading normal estimation research such as [StableNormal](https://github.com/hugoycj/StableNormal) and [GenPercept](https://github.com/aim-uofa/GenPercept). **Your contributions and collaboration push the boundaries of 3D modeling!** """ ) if __name__ == "__main__": # Initialize pipeline pipeline = TrellisImageTo3DPipeline.from_pretrained("Stable-X/trellis-normal-v0-1") pipeline.cuda() # Initialize normal predictor normal_predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal_turbo", trust_repo=True, yoso_version='yoso-normal-v1-8-1') # Launch the app demo.launch(share=True)