app.py

import os
# Force attention backends compatible with the new ZeroGPU (Blackwell) stack.
# Must be set BEFORE any trellis / dinov2 import.
# Trellis's dense attention has a native SDPA path; use it.
os.environ.setdefault('ATTN_BACKEND', 'sdpa')
# Sparse attention only knows 'xformers' or 'flash_attn'; keep 'xformers' but
# monkey-patch xformers.ops.memory_efficient_attention to SDPA below (none of
# the prebuilt xformers ops support sm_120 / Blackwell).
os.environ.setdefault('SPARSE_ATTN_BACKEND', 'xformers')
os.environ.setdefault('SPCONV_ALGO', 'native')
# Force dinov2 (loaded via torch.hub for image conditioning) to take its pure
# torch.nn.functional.scaled_dot_product_attention path instead of importing
# xformers.ops.memory_efficient_attention (which raises on sm_120).
os.environ.setdefault('XFORMERS_DISABLED', '1')

import sys
import subprocess
import tempfile
import ctypes

import spaces
import torch
import gradio as gr

# ---------------------------------------------------------------------------
# xformers -> SDPA shim for Blackwell (sm_120).
# The prebuilt xformers wheel ships FA3, FA2 and CutlassF ops that all assert
# device capability <= (9, 0); none load on sm_120, so any call to
# memory_efficient_attention raises NotImplementedError. dinov2 (image
# conditioning model in trellis) and trellis's own sparse paths both call it.
# Replace memory_efficient_attention with an SDPA-backed implementation that
# also handles xformers.fmha.BlockDiagonalMask (used by sparse attention).
# ---------------------------------------------------------------------------
try:
    import xformers  # noqa: F401
    import xformers.ops as _xops
    from torch.nn.functional import scaled_dot_product_attention as _sdpa

    try:
        _BlockDiagonalMask = _xops.fmha.BlockDiagonalMask
    except Exception:
        _BlockDiagonalMask = None

    def _mea_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, *args, **kwargs):
        # q, k, v: [B, N, H, C] (xformers layout). SDPA expects [B, H, N, C].
        if attn_bias is None:
            qh = q.transpose(1, 2)
            kh = k.transpose(1, 2)
            vh = v.transpose(1, 2)
            out = _sdpa(qh, kh, vh, dropout_p=p, scale=scale)
            return out.transpose(1, 2).contiguous()

        if _BlockDiagonalMask is not None and isinstance(attn_bias, _BlockDiagonalMask):
            # BlockDiagonal: q, k, v come as [1, T, H, C] where T is the
            # concatenation of variable-length blocks. Split, apply SDPA per
            # block, concatenate. q and kv can have different seqlens.
            q_info = attn_bias.q_seqinfo
            kv_info = attn_bias.k_seqinfo
            q_starts = q_info.seqstart_py
            kv_starts = kv_info.seqstart_py
            outs = []
            for i in range(len(q_starts) - 1):
                qs, qe = q_starts[i], q_starts[i + 1]
                ks, ke = kv_starts[i], kv_starts[i + 1]
                qi = q[:, qs:qe].transpose(1, 2)
                ki = k[:, ks:ke].transpose(1, 2)
                vi = v[:, ks:ke].transpose(1, 2)
                oi = _sdpa(qi, ki, vi, dropout_p=p, scale=scale)
                outs.append(oi.transpose(1, 2))
            return torch.cat(outs, dim=1).contiguous()

        # Fallback: dense additive bias.
        qh = q.transpose(1, 2)
        kh = k.transpose(1, 2)
        vh = v.transpose(1, 2)
        out = _sdpa(qh, kh, vh, attn_mask=attn_bias, dropout_p=p, scale=scale)
        return out.transpose(1, 2).contiguous()

    _xops.memory_efficient_attention = _mea_sdpa
    print("[xformers-shim] Replaced memory_efficient_attention with SDPA backend (Blackwell sm_120 fallback).")
except Exception as _e:
    print(f"[xformers-shim] Skipped: {_e}")

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(
        """
        <h1 style='text-align: center;'>Hi3DGen: High-fidelity 3D Geometry Generation from Images via Normal Bridging</h1>
        <p style='text-align: center;'>
            <strong>V0.1, Introduced By 
            <a href="https://gaplab.cuhk.edu.cn/" target="_blank">GAP Lab</a> from CUHKSZ and 
            <a href="https://www.nvsgames.cn/" target="_blank">Game-AIGC Team</a> from ByteDance</strong>
        </p>
        """
    )
    
    with gr.Row():
        gr.Markdown("""
                    <p align="center">
                    <a title="Website" href="https://stable-x.github.io/Hi3DGen/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                        <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
                    </a>
                    <a title="arXiv" href="https://stable-x.github.io/Hi3DGen/hi3dgen_paper.pdf" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
                    </a>
                    <a title="Github" href="https://github.com/Stable-X/Hi3DGen" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                        <img src="https://img.shields.io/github/stars/Stable-X/Hi3DGen?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
                    </a>
                    <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                        <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
                    </a>
                    </p>
                    """)

    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("<div style='text-align: center; padding: 40px; font-size: 24px;'>Multiple Images functionality is coming soon!</div>")
                        
            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)