Update app.py

app.py

@@ -1,19 +1,20 @@
 import gradio as gr
-import spaces
-from gradio_litmodel3d import LitModel3D
-
 import os
 import shutil
-os.environ['SPCONV_ALGO'] = 'native'
-from typing import *
-import torch
+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
-import trimesh
-import tempfile
+
+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')
@@ -35,7 +36,7 @@ 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
+        return None, None, None, None
 
     if seed == -1:
         seed = np.random.randint(0, MAX_SEED)
@@ -59,73 +60,43 @@ def generate_3d(image, seed=-1,
     )
     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
-        mask = np.fromfunction(lambda y, x: x > y, (h, w))
-        mask = mask.astype(bool)
+        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'])]
     
-    # 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
-    
     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)
-    
-    return temp_file_path
+    mesh.export(temp_file.name)
+    return temp_file.name
 
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
-    gr.Markdown(
-        """
+    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">
-                        <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">
-                        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
-                    </a>
-                    <a title="Github" href="https://github.com/Stable-X/Hi3DGen" target="_blank">
-                        <img src="https://img.shields.io/github/stars/Stable-X/Hi3DGen?label=GitHub%20%E2%98%85&logo=github&color=C8C">
-                    </a>
-                    <a title="Social" href="https://x.com/ychngji6" target="_blank">
-                        <img src="https://www.obukhov.ai/img/badges/badge-social.svg">
-                    </a>
-                    </p>
-                    """)
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -134,10 +105,8 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                     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")
@@ -148,18 +117,17 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 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 approximately 40MB, may take around 1 minute to load)")
+                        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"],
@@ -167,17 +135,17 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                             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,  
+            image_prompt, seed,
             ss_guidance_strength, ss_sampling_steps,
             slat_guidance_strength, slat_sampling_steps
         ],
@@ -190,10 +158,9 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     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],
@@ -202,7 +169,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         lambda: gr.Button(interactive=True),
         outputs=[download_btn],
     )
-    
+
     examples = gr.Examples(
         examples=[
             f'assets/example_image/{image}'
@@ -211,19 +178,23 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         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)
-        - **Normal Estimation:** We build on top of [StableNormal](https://github.com/hugoycj/StableNormal)
-        """
-    )
+    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__":
-    # 初始化 pipeline，使用 CPU
+    # ✅ 强制使用 CPU
     pipeline = TrellisImageTo3DPipeline.from_pretrained("Stable-X/trellis-normal-v0-1")
-    
-    # 加载 StableNormal 模型（会自动选择 CPU）
-    normal_predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal_turbo", trust_repo=True, yoso_version='yoso-normal-v1-8-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()