app.py

import gradio as gr
import os
import shutil
import tempfile
import datetime
import numpy as np
import torch
import imageio
import trimesh

from PIL import Image
from typing import *
from gradio_litmodel3d import LitModel3D
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import render_utils

os.environ['SPCONV_ALGO'] = 'native'

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)

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

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, 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]
    
    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):
        h, w, c = color_np.shape
        mask = np.fromfunction(lambda y, x: x > y, (h, w)).astype(bool)
        mask = np.stack([mask] * c, axis=-1)
        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'])]
    
    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):
    if not mesh_path:
        return None
    temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False)
    mesh = trimesh.load_mesh(mesh_path)
    mesh.export(temp_file.name)
    return temp_file.name

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():
        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")

        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 approx. 40MB)")
                    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 acknowledge contributions from:
    - [Trellis 3D](https://github.com/microsoft/TRELLIS)
    - [StableNormal](https://github.com/hugoycj/StableNormal)
    """)

if __name__ == "__main__":
    # ✅ 强制使用 CPU
    pipeline = TrellisImageTo3DPipeline.from_pretrained("Stable-X/trellis-normal-v0-1")
    pipeline.to("cpu")  # <-- 强制使用 CPU

    normal_predictor = torch.hub.load(
        "hugoycj/StableNormal",
        "StableNormal_turbo",
        trust_repo=True,
        yoso_version="yoso-normal-v1-8-1"
    )
    normal_predictor.to("cpu")  # <-- 也强制使用 CPU

    demo.launch()